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 #############################

fitter_settings={
'nsteps': 3000,
'nwalkers': 16, 
'burnin': 1000,
'fit_V_CF': True,
'fit_CP': True, #default: True. If False: closure phases not fitted, if True: closure phases fitted
'fit_SED': False, #default: False1
'deredden_SED': True, #default: False
'A_V': 0.4,
'use_corrflux': False, #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': 151,
'scale_mas_fixed': None, # None, #3.781 #[mas/px] default: None
'zoom_size_mas': 20, #default: None
'draw_rim': False, #default: True, draw rim for the flared disk model
'pk_version': 2,
'best_fit':'median', #'median' (default), or 'mode'
'nodata': False, #default: False
'color_scheme':'wl'
}

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

# L band DX Cha
basedir = '/allegro6/matisse/varga/targets3/DX_Cha/' 
inputdir = basedir
outputdir = basedir+'/modelling/' 
oifits_paths = [#inputdir+'/DX_Cha_2020-03-22T05_54_46_L_TARGET_CHOPPED_FINALCAL_INT.fits',
inputdir+'DX_Cha_2022-01-23T07_37_08_L_TARGET_FINALCAL_INT.fits',
inputdir+'DX_Cha_2022-01-23T07_37_08_N_TARGET_FINALCAL_INT_VISCORRFLUX.fits'
#inputdir+'/DX_Cha_2021-03-11T04_45_57_L_TARGET_CHOPPED_FINALCAL_INT.fits',
#inputdir+'/DX_Cha_2022-01-23T07_37_08_L_TARGET_FINALCAL_INT.fits',
]
oifits_datatypes = [
['int','t3'],
['int','t3'],
]
instrument = 'MATISSE' #MIDI
correct_errors = False
error_terms = {}
SEDs = []

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

#available_models: 
#   'gaussian'
#   'ring_asymm
#   'ring_symm'
#   'temp_grad_asymm' 
#   'temp_grad_symm' 
#   'temp_grad_flared_symm'
#   'cg_min_symm' 
#   'cg_min_asymm' 
#   'radimg_symm'
#   'radimg_asymm'
#   'central_1D_symm'
#   'central_1D_asymm'
#   'temp_grad_1D_symm'
#   'temp_grad_1D_asymm'
#   'banana'

model_name = 'temp_grad_1D_asymm' #'temp_grad_asymm'   #'ring_asymm'  #'temp_grad_flared_symm' #
run_name = 'LN_temp_grad_1D_asymm_plus_2stars'
wl_dic = {
'wavelengths' : [3.188707,8.60937],  #(um) [3.2,3.45,3.7]
'bandwidths' : [0.2,0.4], #(um)
#'wl_min' : 1.59, #1.59, #3.1, #3.1, #8.0, #3.3, #0.2
#'wl_max' : 8.4, #12.0, #8.5, #12.5, #11.0, #13.0, #10000.0
}

###################### Initial parameter values ########################
L_star =  57.8 #L_Sun GAIA DR3
dist_pc = 102.4 #GAIA DR3
position_angle_init = (122.6)*math.pi/180.0  
axis_ratio_init = 0.82
modulation_amplitude_init = 0.3 
modulation_angle_init = 307.1*math.pi/180.0
r_mod_max_init = 0.0
sigma_r_mod_init = 0.0
xc_init = 1.0 #0.8
yc_init = 1.0 #-0.8
f_star_init = 0.2 #-0.1
#total_flux_star_init = 0.0 #0.13 #Jy - only in correlated flux mode
#total_flux_init = 0.0 #3.4 #20.8 # Jy (if use_corrflux == True)
alpha_f_star_init = 2.0 #spectral index of the stellar flux fraction
V_background_init = 0.0 #0.52# 0.27 #0.27 #0.42
# Gaussian
#FWHM_init = 9.18 #(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)
# temp_grad 
T_in_init = 1500.0
r_in_mas_init = 3.0 #5.0
q_init = 0.76
#r_gap_in_mas_init = 0.58
#r_gap_out_mas_init = 15.08
# extra point source
f_comp_init = 0.2
alpha_f_comp_init = 2.0 
xc_comp_init = -1.0 #1.16 #1.44 
yc_comp_init = -1.0 #0.85 # 0.00
#2021:
'''
f_star_init = 0.0
f_comp_init = 0.08 
xc_comp_init = -0.48
yc_comp_init = -2.9
FWHM_init = 9.82 #(mas)
'''
############################# 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
# geometry and flux parameters
parameter_dic = {
'common': [
modelparameter('position_angle',      F ,[0.0,1.0*np.pi],position_angle_init,None,r'\theta',r'\degree',r'\mathrm{rad}','%.1f',180.0/np.pi),
modelparameter('axis_ratio',          F ,[0.25,1.0],axis_ratio_init,None,r'\cos\,i','',None,'%.2f',1.0),
modelparameter('xc',                  T ,[-7.0,7.0],xc_init,[xc_init-1.0,xc_init+1.0],r'x_c','\mathrm{mas}',None,'%.2f',1.0),
modelparameter('yc',                  T ,[-7.0,7.0],yc_init,[yc_init-1.0,yc_init+1.0],r'y_c','\mathrm{mas}',None,'%.2f',1.0),
modelparameter('f_star',              T ,[0.0,1.0],f_star_init,[f_star_init*0.95,f_star_init*1.5],'f_\mathrm{star}','',None,'%.2f',1.0),
modelparameter('alpha_f_star',        T ,[0.0,3.0],alpha_f_star_init,None,r'\alpha_{f\mathrm{star}}','',None,'%.2f',1.0),
#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),
#modelparameter('total_flux',          F ,[0.01,150.0],total_flux_init,None,r'F_\mathrm{tot}','\mathrm{Jy}',None,'%.1f',1.0),
modelparameter('V_background',        F ,[0.0,1.0],V_background_init,None,r'V_\mathrm{bg}','',None,'%.2f',1.0),
modelparameter('dist_pc',             F ,[50.0,1000.0],dist_pc,None,r'd',r'\mathrm{pc}',None,'%.1f',1.0),
#modelparameter('log_f',              F ,[-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 asymmetry
'common_asymm': [
modelparameter('modulation_amplitude',T,[0.0,1.0],modulation_amplitude_init,None,r'A_\mathrm{mod}','',None,'%.2f',1.0),
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),
modelparameter('modulation_radius',   F ,[0.01,30.0],r_mod_max_init,None,r'r_\mathrm{mod,max}','\mathrm{mas}',None,'%.2f',1.0),
modelparameter('modulation_radial_width',F ,[0.01,100.0],sigma_r_mod_init,None,r'\sigma_\mathrm{r,mod}','\mathrm{mas}',None,'%.2f',1.0)
],
# parameters of the temperature gradient model
'temp_grad': [
modelparameter('L_star',              F ,[0.1,100.0],L_star,None,r'L_*',r'\mathrm{L}_\odot',None,'%.1f',1.0),
modelparameter('r_in_mas',            T ,[0.1,150.0],r_in_mas_init,None,r'R_\mathrm{in}','\mathrm{mas}',None,'%.2f',1.0),
modelparameter('q',                   T ,[0.3,5.0],q_init,None,r'q','',None,'%.2f',1.0),
modelparameter('T_in',                F ,[0.01,2000.0],T_in_init,None,r'T_\mathrm{in}','\mathrm{K}',None,'%.1f',1.0),
#modelparameter('r_gap_in_mas',        T ,[0.1,150.0],r_gap_in_mas_init,None,r'r_\mathrm{gap,in}','\mathrm{mas}',None,'%.2f',1.0),
#modelparameter('r_gap_out_mas',       T ,[0.1,30.0],r_gap_out_mas_init,None,r'r_\mathrm{gap,out}','\mathrm{mas}',None,'%.2f',1.0)
],
#'gaussian': [
#modelparameter('FWHM',                T ,[0.1,50.0],FWHM_init,None,r'FWHM','\mathrm{mas}',None,'%.2f',1.0),
#modelparameter('Lor_frac',            F, [0.0,1.0],Lor_frac_init, None, r'Lor','',None,'%.2f',1.0),
#],
'extra_point_source' : [
modelparameter('f_comp',              T, [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),
modelparameter('alpha_f_comp',        T ,[0.0,3.0],alpha_f_comp_init,None,r'\alpha_{f\mathrm{comp}}','',None,'%.2f',1.0),
modelparameter('xc_comp',             T ,[-7.0,7.0],xc_comp_init,[xc_comp_init-1.0,xc_comp_init+1.0],r'x_{\mathrm{comp}}','\mathrm{mas}',None,'%.2f',1.0),
modelparameter('yc_comp',             T ,[-7.0,7.0],yc_comp_init,[yc_comp_init-1.0,yc_comp_init+1.0],r'y_{\mathrm{comp}}','\mathrm{mas}',None,'%.2f',1.0),
],
}

model_options = {'nu0':c0/(wl_dic['wavelengths'][0]/1e6),'nu':np.nan,'nx':fitter_settings['img_size_px'],'ring_width_ratio':0.067,
    'scale_mas':None, 'calc_profile':False,'wl_dic':wl_dic,
    'V_CF_chi2_weight':1.0,'wl_min':3.0,'wl_max':13.0,
    #'posterior_range_limits' : [{'param_key':'r_in','min':0.1,'max':math.inf},
    #                            {'param_key':'tau_mp','min':1e-8,'max':1.0},
    #                            {'param_key':'C_i','min':1e-8,'max':0.002}]
    #                            {'param_key':'C_o','min':1e-8,'max':0.0004}
    #                           ]
                            #{'param_key':'r_qjump','min':0.001,'max':2.0},
    #                    [{'param_key':'c_fos_i','min':-math.inf,'max':100.2},
    #                      {'param_key':'c_ens_i','min':-math.inf,'max':2.0},
    #                      {'param_key':'c_enl_i','min':-math.inf,'max':2.0},
    #                      ]
    }