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 = True #default: True. If False: closure phases not fitted, if True: closure phases fitted
use_corrflux = False #default: False. If True: fit correlated fluxes, if False: fit visibilities
use_chi2 = True #default: ?, if False: use likelihood (with parameter log f) for the fitting
img_size_px = 301 #151
scale_mas_fixed = None #3.781 #[mas/px] default: None
zoom_size_mas = 10.0 #default: None
draw_rim = False #default: True, draw rim for the flared disk model

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

# L band HD 163296
basedir = '/allegro6/matisse/varga/targets2/HD_163296/' #'/data2/archive/data/MATISSE/targets2/HD_163296/' #'/allegro6/matisse/varga/targets/HD_163296/' #'/data2/archive/data/MATISSE/targets/HD_163296'
inputdir = basedir
outputdir = basedir+'/modelling/visibility_modelling/3um3/new' #'modelling/visibility_modelling/3um9/'
#'/data2/archive/data/MATISSE/targets2/HD_163296/' #'/allegro6/matisse/varga/targets/HD_163296/' 
#'/data2/archive/data/MATISSE/targets/HD_163296'
oifits_paths = [inputdir+'/HD_163296_2019-03-23T08_41_19_L_TARGET_FINALCAL_INT.fits',
 inputdir+'/HD_163296_2019-05-06T08_19_51_L_TARGET_FINALCAL_INT.fits',
 #inputdir+'/HD_163296_2019-06-26T06_26_41_L_TARGET_CHOPPED_FINALCAL_INT.fits',
 #inputdir+'/HD_163296_2019-06-26T06_26_41_L_TARGET_FINALCAL_INT.fits',
 #inputdir+'/HD_163296_2019-06-29T07_07_44_L_TARGET_CHOPPED_FINALCAL_INT.fits',
 #inputdir+'/HD_163296_2019-06-29T07_07_44_L_TARGET_FINALCAL_INT.fits',
#inputdir+'/Aspro2_HD_163296_MATISSE_LM_2_86542-4_18239-62ch_A0-G1-J2-J3_2020-06-10.fits'
]
instrument = 'MATISSE' #MIDI
correct_errors = True
V_err_lower_limit = 0.03
CP_err_lower_limit = 1.0*math.pi/180.0

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

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

model_name = 'temp_grad_asymm'   #'gaussian' #'ring_asymm'  #'temp_grad_flared_symm' #
run_name = 'temp_grad_asymm_errors_equalized_multiwl_chi2' 
wavelengths = [3.2,3.45,3.7] #[3.3] #[3.15,3.5,3.95] #[3.5] # #[3.3] #[3.1,3.5,3.9] #(um)
bandwidths = [0.2,0.2,0.2] #(um)

#'temp_grad_asymm_errors_equalized_multiwl_chi2_2019May_fitasymm' 

###################### Initial parameter values ########################
L_star =  16.0 #40.0 L_Sun
dist_pc = 101.2 #183.9 # 
position_angle_init = (143.0)*math.pi/180.0  #150 (313.0-180.0)*math.pi/180.0 #ALMA: 313 deg #133.0*math.pi/180.0  #140.9 130.4*math.pi/180.0 #140.9
axis_ratio_init = 0.56 #0.8 0.68 #ALMA: 0.68 #0.82 #0.68#0.66 #0.65 
modulation_amplitude_init = 0.54 #1.0 #0.61 #0.51 #0.0
modulation_angle_init = 182.0*math.pi/180.0 # 173.1*math.pi/180.0 #210.3
xc_init = 0.0 #0.8
yc_init = 0.0 #-0.8
f_disk_init = 1.0-0.066 #-0.1
#HD 97048: @3.58 um (cont). total flux 3.2 Jy, stellar flux 0.32 Jy
#-0.074948 #0.066 #-0.066 # flux fraction of the disk component (fixed: 0.77/11.6 Jy = 0.066 at 3.3 um)
total_flux_star_init = 0.0 #0.13 #Jy - only in correlated flux mode
total_flux_init = 4.8 #3.4 #20.8 # Jy (if use_corrflux == True)
alpha_f_star_init = 1.6842 #spectral index of the stellar flux fraction
V_background_init = 0.0 #0.52# 0.27 #0.27 #0.42
log_f_init = -1.0 #-1.92# -1.82 #-2.94
# 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 = 1.71 #66.61
q_init = 0.66 #0.77
# 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
# Gaussian
FWHM_init = 6.51 #3.76 #2.44 #16.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, 1 for Lorentzian)
# ring
ring_radius_init = 2.76 #66.61 #3.0 #0.2 #0.55 #2.77 #star + offset Gaussian: 0.55 #mas
kernel_width_FWHM_init = 0.0 #4.99 #8.0 #5.08  #6.07 #2.9 #2.9 #6.0 #mas
#extra Gaussian
f_extra_Gaussian_init = 0.0 #0.11 #0.76 #0.58  #5 #0.69 #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',      T ,pa_range,position_angle_init,None,r'\theta',r'\degree',r'\mathrm{rad}','%.1f',180.0/np.pi)
axra = modelparameter('axis_ratio',          T ,[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',          F ,[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',T ,[0.0,4.0],modulation_amplitude_init,None,r'A_\mathrm{mod}','',None,'%.2f',1.0)
mang = modelparameter('modulation_angle',    T ,[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',                F ,[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',            T ,[0.1,150.0],r_in_mas_init,None,r'R_\mathrm{in}','\mathrm{mas}',None,'%.2f',1.0)
qqqq = modelparameter('q',                   T ,[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}