import math
import numpy as np
from intmodlib import *

sigma = 5.670373e-8; #(W m-2 K-4) Stefan-Boltzmann constant
T_sub = 1500.0 #K

########################## Fitter settings #############################

nsteps = 10000 #10000
nwalkers = 32 #32
burnin = 1000 #1000
fit_CP = False #default: True. If False: closure phases not fitted, if True: closure phases fitted
use_corrflux = True #default: False. If True: fit correlated fluxes, if False: fit visibilities
use_chi2 = True #default: True, if False: use likelihood (with parameter log f) for the fitting
img_size_px = 201 #301 #201 #151
scale_mas_fixed = None #3.781 #[mas/px] default: None
zoom_size_mas = 30.0 #default: None
draw_rim = False #default: True, draw rim for the flared disk model

########################### Data settings ##############################

# N band HD 163296
basedir = '/allegro6/matisse/varga/targets2/HD_163296/'
inputdir = basedir
outputdir = basedir+'/modelling/N/MIDI_chi2/' 
oifits_paths = [inputdir+'HD_163296_MIDI.fits']
instrument = 'MIDI'

correct_errors = False
CF_error_percent_lower_limit = 0.0
V_err_lower_limit = 0.0
CP_err_lower_limit = 0.0

########################## Model settings #############################

#available_models: 
#   'gaussian'
#   'ring_asymm
#   'ring_symm'
#   'temp_grad_asymm' 
#   'temp_grad_symm' 
#   'temp_grad_flared_symm'

model_name = 'gaussian' #'temp_grad_asymm'   #'ring_asymm'  #'temp_grad_flared_symm' #
run_name = 'MIDI_central_gaussian_chi2'
wavelengths = [9.5] #8.5, 9.5, 10.5, 11.5, 12.5 #(um)
bandwidths = [1.0] #(um)

###################### Initial parameter values ########################
# initial parameters N
L_star = 16.0 #L_Sun
dist_pc = 101.2 
# Gaussian
FWHM_init = 18.0 #7.0 #7.0 #(mas) FWHM of a central Gaussian component (only in temp grad model)
# Gaussian and ring model
Lor_frac_init = 0.0 #Fraction of lorentzian for Gaussian model (0 for Gaussian)
position_angle_init = 133.0*math.pi/180.0 #ALMA
axis_ratio_init = 0.68 #ALMA
f_disk_init = 1.0-0.0 #-0.066 # flux fraction of the disk component (fixed: 0.77/11.6 Jy = 0.066 at 3.3 um)
f_uniform_init = 0.0
log_f_init = -9.29 #-1.92# -1.82 #-2.94
total_flux_init = 19.9 #Jy 12.6,19.9,20.4,16.6,10.7 @8.5,9.5,10.5,11.5,12.5um  (if use_corrflux == True)
total_flux_star_init = 0.11 #Jy 0.13,0.11,0.09,0.07,0.06 @8.5,9.5,10.5,11.5,12.5um - only in correlated flux mode
alpha_f_star_init = 1.6842 #spectral index of the stellar flux fraction
V_background_init = 0.0

modulation_amplitude_init = 0.0
modulation_angle_init = 173.1*math.pi/180.0 # 173.1*math.pi/180.0 #210.3
xc_init = 0.0 #0.8
yc_init = 0.0 #-0.8
log_f_init = -1.0
# temp_grad
r_in_mas_init = np.sqrt(L_star*3.846e26/(4.0*np.pi*sigma*(T_sub**4)))/1.4960e11/dist_pc*1000.0 #set to the sublimation radius
r_in_mas_init = 2.6 #66.61
q_init = 0.69
# flared disk
rim_width_mas_init = 0.0 #2.86
p_rim_init = 1.0
surface_angle_init = 7.7*math.pi/180.0 #14.5 #31 deg
# ring
ring_radius_init = 2.76 
kernel_width_FWHM_init = 0.0 
#extra Gaussian
f_extra_Gaussian_init = 0.0 #flux fraction of a central Gaussian component, with respect to the disk flux
FWHM_extra_Gaussian_init = 0.17
position_angle_extra_Gaussian_init = (0.0)*math.pi/180.0  
axis_ratio_extra_Gaussian_init = 1.0 
xc_extra_Gaussian_init = -1.47
yc_extra_Gaussian_init = 1.22
# extra ring
ring_radius_extra_init = 0.0 #1.76
f_ring_extra_init = 0.0 #0.42
alpha_f_ring_extra_init = 0.0 #spectral index of the extra ring
# extra uniform disk
f_uniform_init = 0.0
# extra point source
f_comp_init = 0.0 #0.1
xc_comp_init = -1.45 #-368.0
yc_comp_init = 1.18 #258.0

############################# Model setup ##############################

# - define models with parameters
if model_name.startswith('temp_grad_flared'):
    pa_range = [0.0,2.0*np.pi]
else:
    pa_range = [0.0,1.0*np.pi]

F = False
T = True
# arguments of modelparameter init function: param_key,fitted,priors,init_value,param_range,label,unit,print_format,print_factor
# common parameters
posa = modelparameter('position_angle',      F ,pa_range,position_angle_init,None,r'\theta',r'\degree',r'\mathrm{rad}','%.1f',180.0/np.pi)
axra = modelparameter('axis_ratio',          F ,[0.15,1.0],axis_ratio_init,None,r'\cos\,i','',None,'%.2f',1.0)
xcen = modelparameter('xc',                  F ,[-7.0,7.0],xc_init,[xc_init-1.0,xc_init+1.0],r'x_c','\mathrm{mas}',None,'%.2f',1.0)
ycen = modelparameter('yc',                  F ,[-7.0,7.0],yc_init,[yc_init-1.0,yc_init+1.0],r'y_c','\mathrm{mas}',None,'%.2f',1.0)
fdsk = modelparameter('f_disk',              F ,[0.0,1.0],f_disk_init,[f_disk_init*0.95,f_disk_init*1.01],'f_\mathrm{disk}','',None,'%.2f',1.0)
alst = modelparameter('alpha_f_star',        F ,[0.0,3.0],alpha_f_star_init,None,r'\alpha_{f\mathrm{star}}','',None,'%.2f',1.0)
tofs = modelparameter('total_flux_star',     F ,[0.01,100.0],total_flux_star_init,None,r'F_\mathrm{tot,star}','\mathrm{Jy}',None,'%.2f',1.0)
tofl = modelparameter('total_flux',          T ,[0.01,150.0],total_flux_init,None,r'F_\mathrm{tot}','\mathrm{Jy}',None,'%.1f',1.0)
vbgr = modelparameter('V_background',        F ,[0.0,1.0],V_background_init,None,r'V_\mathrm{bg}','',None,'%.2f',1.0)
dipc = modelparameter('dist_pc',             F ,[50.0,1000.0],dist_pc,None,r'd',r'\mathrm{pc}',None,'%.1f',1.0)
logf = modelparameter('log_f',               T ,[-10.0,2.0],log_f_init,[log_f_init-1.0,log_f_init+1.0],r'\log\,f','',None,'%.2f',1.0)
# parameters of the asymmetric ring and the asymmetric temperature gradient model
mamp = modelparameter('modulation_amplitude',F ,[0.0,4.0],modulation_amplitude_init,None,r'A_\mathrm{mod}','',None,'%.2f',1.0)
mang = modelparameter('modulation_angle',    F ,[0.0,2.0*np.pi],modulation_angle_init,None,r'\phi_\mathrm{mod}',r'\degree',r'\mathrm{rad}','%.1f',180.0/np.pi)
# parameter of the Gaussian model
fwhm = modelparameter('FWHM',                T ,[0.1,50.0],FWHM_init,None,r'FWHM','\mathrm{mas}',None,'%.2f',1.0)
# parameters of the smoothed ring model
rngr = modelparameter('ring_radius',         F ,[0.1,100.0],ring_radius_init,None,r'R_\mathrm{in}','\mathrm{mas}',None,'%.2f',1.0)
kewi = modelparameter('kernel_width_FWHM',   F ,[0.0,150.0],kernel_width_FWHM_init,None,r'FWHM_\mathrm{kernel}','\mathrm{mas}',None,'%.2f',1.0)
# parameter of the Gaussian model and of the ring model
lorf = modelparameter('Lor_frac',            F, [0.0,1.0],Lor_frac_init, None, r'Lor','',None,'%.2f',1.0)
# parameters of the temperature gradient model and of the flared disk model
lust = modelparameter('L_star',              F ,[0.1,100.0],L_star,None,r'L_*',r'\mathrm{L}_\odot',None,'%.1f',1.0)
rinm = modelparameter('r_in_mas',            F ,[0.1,150.0],r_in_mas_init,None,r'R_\mathrm{in}','\mathrm{mas}',None,'%.2f',1.0)
qqqq = modelparameter('q',                   F ,[0.3,5.0],q_init,None,r'q','',None,'%.2f',1.0)
# parameters of the flared disk model
suan = modelparameter('surface_angle',       F, [0.0,40.0*np.pi/180.0],surface_angle_init,[surface_angle_init*0.95,surface_angle_init*1.01],r'\psi',r'\degree',r'\mathrm{rad}','%.1f',180.0/np.pi)
rimw = modelparameter('rim_width_mas',       F ,[0.0,10.0],rim_width_mas_init,None,r'w_\mathrm{rim}','\mathrm{mas}',None,'%.2f',1.0)
prim = modelparameter('p_rim',               F ,[0.001,10.0],p_rim_init,None,r'p_\mathrm{rim}','',None,'%.2f',1.0)
# parameters for additional components
# extra Gaussian
fega = modelparameter('f_extra_Gaussian',    F, [0.0,1.0],f_extra_Gaussian_init,[f_extra_Gaussian_init*0.95,f_extra_Gaussian_init*1.01],r'f_\mathrm{blob}','',None,'%.2f',1.0)
fweg = modelparameter('FWHM_extra_Gaussian', F ,[0.0001,50.0],FWHM_extra_Gaussian_init,None,r'FWHM_\mathrm{blob}','\mathrm{mas}',None,'%.2f',1.0)
paeg = modelparameter('position_angle_extra_Gaussian',F ,pa_range,position_angle_extra_Gaussian_init,None,r'\theta_\mathrm{blob}',r'\degree',r'\mathrm{rad}','%.1f',180.0/np.pi)
areg = modelparameter('axis_ratio_extra_Gaussian',    F ,[0.15,1.0],axis_ratio_extra_Gaussian_init,None,r'\cos\,i_\mathrm{blob}','',None,'%.2f',1.0)
xceg = modelparameter('xc_extra_Gaussian',   F ,[-74.0,74.0],xc_extra_Gaussian_init,[xc_extra_Gaussian_init-1.0,xc_extra_Gaussian_init+1.0],r'x_{\mathrm{blob}}','\mathrm{mas}',None,'%.2f',1.0)
yceg = modelparameter('yc_extra_Gaussian',   F ,[-74.0,74.0],yc_extra_Gaussian_init,[yc_extra_Gaussian_init-1.0,yc_extra_Gaussian_init+1.0],r'y_{\mathrm{blob}}','\mathrm{mas}',None,'%.2f',1.0)
# extra uniform disk
funi = modelparameter('f_uniform',           F, [0.0,1.0],f_uniform_init,[f_uniform_init*0.95,f_uniform_init*1.01],r'f_\mathrm{uniform}','',None,'%.2f',1.0)
# extra ring
feri = modelparameter('f_ring_extra',        F ,[0.0,1.0],f_ring_extra_init,[f_ring_extra_init-0.1,f_ring_extra_init+0.1],'f_\mathrm{ring2}','',None,'%.2f',1.0)
rrer = modelparameter('ring_radius_extra',   F ,[0.1,10.0],ring_radius_extra_init,None,r'R_\mathrm{ring2}','\mathrm{mas}',None,'%.2f',1.0)
afre = modelparameter('alpha_f_ring_extra',  F, [-4.0,4.0],alpha_f_ring_extra_init,[alpha_f_ring_extra_init-0.5,alpha_f_ring_extra_init+0.5],r'\alpha_{f\mathrm{ring2}}','',None,'%.2f',1.0)
# extra point source
feps = modelparameter('f_comp',              F, [0.0,1.0],f_comp_init,[f_comp_init*0.95,f_comp_init*1.01],r'f_\mathrm{companion}','',None,'%.2f',1.0)
xcep = modelparameter('xc_comp',             F ,[-74.0,74.0],xc_comp_init,[xc_comp_init-1.0,xc_comp_init+1.0],r'x_{\mathrm{comp}}','\mathrm{mas}',None,'%.2f',1.0)
ycep = modelparameter('yc_comp',             F ,[-74.0,74.0],yc_comp_init,[yc_comp_init-1.0,yc_comp_init+1.0],r'y_{\mathrm{comp}}','\mathrm{mas}',None,'%.2f',1.0)

#available_models: 'gaussian','ring_asymm, 'ring_symm', 'temp_grad_asymm', 'temp_grad_symm' , 'temp_grad_flared_symm' 
img_options = {'nu0':c0/(wavelengths[0]/1e6),'nu':np.nan,'nx':img_size_px,'ring_width_ratio':0.1,'scale_mas':None, 'calc_profile':False}

########################### Misc. settings #############################

# limit posterior range, list of dicts
#posterior_range_limits = [{'param_key':'FWHM','min':-math.inf,'max':30.0},
#                          ]