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 = 3000 #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: True, if False: use likelihood (with parameter log f) for the fitting
img_size_px = 201 #301 
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 ##############################

# H band HD 163296 PIONIER
basedir = '/allegro6/matisse/varga/targets2/HD_163296/' 
inputdir = basedir + 'PIONIER/'
outputdir = basedir+'/modelling/visibility_modelling/H_band/' 
oifits_paths = [
#inputdir+'/PIONIER.2013-06-15T05 16 46_oidataCalibrated.fits',
#inputdir+'/PIONIER.2013-07-03T05 31 08_oidataCalibrated.fits',
#inputdir+'/PIONI.2014-06-24T07 50 56.392_oidataCalibrated.fits',
#inputdir+'/PIONI.2014-06-24T08 13 12.448_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-04-10T06 00 55.004_oidataCalibrated.fits',
# inputdir+'/PIONI.2017-04-11T07 50 41.473_oidataCalibrated.fits',
# inputdir+'/PIONI.2017-04-11T08 26 11.683_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-04-11T08 49 23.371_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-04-12T09 37 38.090_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-04-22T08 00 20.291_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-04-22T08 30 44.901_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-04-22T09 06 32.011_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-04-22T09 36 55.742_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-05-19T09 12 28.761_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-06-18T01 33 26.901_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-06-19T05 45 50.865_oidataCalibrated.fits',
#inputdir+'/PIONI.2017-06-19T06 56 22.099_oidataCalibrated.fits',
#inputdir+'/PIONI.2018-04-08T08 39 34.027_oidataCalibrated.fits',
inputdir+'/PIONI.2018-06-06T07 12 45.370_oidataCalibrated.fits',
inputdir+'/PIONI.2018-06-06T07 39 48.782_oidataCalibrated.fits',
inputdir+'/PIONI.2018-06-06T08 07 29.166_oidataCalibrated.fits',
#inputdir+'/PIONI.2018-06-06T08 37 57.489_oidataCalibrated.fits',
inputdir+'/PIONI.2018-06-07T05 06 13.830_oidataCalibrated.fits',
inputdir+'/PIONI.2018-06-07T05 50 28.888_oidataCalibrated.fits',
# inputdir+'/PIONI.2019-06-09T05 55 29.359_oidataCalibrated.fits',
# inputdir+'/PIONI.2019-06-09T07 04 43.697_oidataCalibrated.fits',
# inputdir+'/PIONI.2019-06-09T07 38 28.742_oidataCalibrated.fits',
# inputdir+'/PIONI.2019-06-09T08 11 35.580_oidataCalibrated.fits',
#inputdir+'/PIONI.2018-06-19T07 26 49.190_oidataCalibrated.fits',
#inputdir+'/PIONI.2018-06-19T08 27 26.030_oidataCalibrated.fits',
#inputdir+'/PIONI.2019-07-10T03 15 55.994_oidataCalibrated.fits',
#inputdir+'/PIONI.2019-07-10T03 35 19.858_oidataCalibrated.fits',
#inputdir+'/PIONI.2019-07-10T06 56 34.603_oidataCalibrated.fits',
#inputdir+'/PIONI.2019-07-20T01 49 23.278_oidataCalibrated.fits',
#inputdir+'/PIONI.2019-07-20T02 03 42.503_oidataCalibrated.fits',
#inputdir+'/PIONI.2019-07-20T03 35 41.328_oidataCalibrated.fits',
#inputdir+'/PIONI.2019-07-20T03 50 33.494_oidataCalibrated.fits',
#inputdir+'/PIONI.2019-07-29T23 15 39.222_oidataCalibrated.fits',
]
instrument = 'PIONIER'
correct_errors = False
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 = 'ring_asymm'  #'gaussian' #'temp_grad_asymm'   #''temp_grad_flared_symm' #
run_name = 'ring_asym_conv_Lorentz_PIONIER_2018Jun_dipper_chi2'
wavelengths = [1.57,1.7] #[1.615,1.69,1.77] # #(um)
bandwidths = [0.07,0.07] #[0.04,0.04,0.04] #(um)

###################### Initial parameter values ########################
L_star =  16.0 #40.0 L_Sun
dist_pc = 101.2 #183.9 # 
position_angle_init = (128.1)*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.62 #0.8 0.68 #ALMA: 0.68 #0.82 #0.68#0.66 #0.65 
modulation_amplitude_init = 0.66 #1.0 #0.61 #0.51 #0.0
modulation_angle_init = 185.4*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.39 #(at 1.615 um) # flux fraction of the disk component (fixed: 2.5/6.48 Jy = 0.39 at 1.615 um)
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.271 #H band #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.73 #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 = 1.0 #Fraction of lorentzian for Gaussian model (0 for Gaussian, 1 for pseudo-Lorentzian)
# ring
ring_radius_init = 2.16 #66.61 #3.0 #0.2 #0.55 #2.77 #star + offset Gaussian: 0.55 #mas
kernel_width_FWHM_init = 4.55 #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

# initial parameters N
# L_star = 16.0 #L_Sun
# dist_pc = 101.2 
# 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
# q_init = 0.69 #0.7
# FWHM_init = 16.0 #7.0 #7.0 #(mas) FWHM of a central Gaussian component (only in temp grad model)
# kernel_width_FWHM_init = 0.0 #2.9 #2.9 #6.0 #mas
# position_angle_init = 133.0*math.pi/180.0 #122.0*math.pi/180.0  #133.0*math.pi/180.0  #140.9 130.4*math.pi/180.0 #140.9
# axis_ratio_init = 0.68 #0.82 #0.68#0.66 #0.65 
# 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 = 12.64 #17.5 # Jy (if use_corrflux == True)
# V_background_init = 0.0

# flux_correction = [1.0,18.5/17.7,18.5/21.9]

############################# 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',              T ,[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',         T ,[0.1,100.0],ring_radius_init,None,r'R_\mathrm{in}','\mathrm{mas}',None,'%.2f',1.0)
kewi = modelparameter('kernel_width_FWHM',   T ,[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}