import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from astropy.io import fits as pyfits
from scipy.interpolate import interp1d
import sys
from matplotlib.colors import LogNorm
from astropy.table import Table
from matplotlib.backends.backend_pdf import PdfPages
from astropy.utils.data import get_pkg_data_filename
from astropy.modeling import models
from scipy import interpolate,integrate
from scipy.optimize import fsolve
from astropy.modeling import models

# TGMdust model

# new model!
# assuming optically thin dust
# the dust temperature is consistently calculated from the stellar SED and from the dust opacity
# the dust column density is assumed to be a power law
# a gap can be included

#Use matplotlib version 1.5.3

show_plots = False

# some constants:
#h = 6.6262e-27
#c = 2.997925e+10 # in cm/s
#au = 1.495978707e13 # in cm
#kboltz = 1.38e-16
#rsun = 6.9634e10 # in cm
#parsec = 3.0857e18 # in cm
limit = 1.0e-20

#def bbody(freq,t):
#    expon = (h*freq)/(kboltz*t)
#    return (2*h*(freq**3)/(c*c))*(1.0/(np.exp(expon) - 1.0))

h = 6.62606957e-34; #(J s)
kboltz = 1.3806488e-23; #(J K-1)
c0 = 299792458.0; #(m s-1)
au = 1.495978707e11 # (m)
rsun = 6.9634e8 # (m)
parsec = 3.0857e16 # (m)

#Planck's law
def B_nu(T,nu):
    return 2.0*h*nu**3.0/(c0**2.0)/(np.exp(h*nu/(kboltz*T))-1)  #[W sr-1 m-2 Hz-1]

# we use 5 dust species:
# amorphous Mg2SiO4 with optical constants from Jaeger et al. (2003)
# amorphous MgSiO3 with optical constants from Dorschner et al. (1995)
# amorphous SiO2 with optical contstants from Henning & Mutschke (1997)
# forsterite Mg2SiO4 at T = 295 K from Suto et al. (2006)
# enstatite MgSiO3 from Zeidler et al. (2013) with T = 300 K
# for each dust species we use 2 sizes : 0.1 and 2 microns
# so we have 10 dust components

def calc_dust_masses(model):
    # disk parameters
    rtaper= model.parameters['r_taper_au'].value # in AU
    #t_subl = model.parameters['T_sub'].value # maximum dust temperature in K
    Nzones = model.model_options['Nzones']
    r_in_lst = []
    r_out_lst = []
    p_lst = []
    n_lst = []
    for i in range(Nzones):
        r_in_lst.append(model.parameters['r_i%d'%(i+1)].value)
        p_lst.append(model.parameters['p%d'%(i+1)].value)
        n_lst.append(10.0**model.parameters['n%d'%(i+1)].value)
        if i>0:
            r_out_lst.append(r_in_lst[-1])
    r_out_lst.append(rtaper) 
    r_in = r_in_lst[0]
    if 'kappa0' in model.parameters:
        kappa0 = model.parameters['kappa0'].value
    Nc = len(model.model_options['cs'])
 
    if 'kappa0' in model.parameters:
        pki = model.model_options['parkappa_idx']
        for i in range(Nzones):
            lambda_break = model.parameters['lambda_break_z%d'%(i+1)].value
            p0_kappa = model.parameters['p0_kappa_z%d'%(i+1)].value
            p1_kappa = model.parameters['p1_kappa_z%d'%(i+1)].value
            f = models.SmoothlyBrokenPowerLaw1D(amplitude=kappa0, x_break=lambda_break,
                                        alpha_1=p0_kappa, alpha_2=p1_kappa,delta=0.2)
            model.parameters['sdens_z%d_%d'%((i+1),pki)].q = f(model.parameters['sdens_z%d_%d'%((i+1),pki)].wl_q) 
            model.parameters['sdens_z%d_%d'%((i+1),pki)].sed_q = f(model.parameters['sdens_z%d_%d'%((i+1),pki)].wl_sed_q)
            model.model_options['temp_cal_pm'][pki+i*int(Nc/Nzones)]['opac_orig'] = f(model.model_options['temp_cal_pm'][pki+i*int(Nc/Nzones)]['wl_orig'])
            idx = model.model_options['cs'][pki+i*int(Nc/Nzones)]
            model.parameters[idx].q = f(model.parameters[idx].wl_q)
            model.parameters[idx].sed_q = f(model.parameters[idx].wl_sed_q)    
    kappas = []
    for i in range(Nc):
        kappas.append(model.parameters[model.model_options['cs'][i]].q)
    rinner = r_in
    # fill the radius grid
    radius = np.logspace(np.log10(rinner),np.log10(rtaper),model.model_options['nx'], endpoint='true')

    surfdens_profiles = []
    ps = np.repeat(p_lst,(int(Nc/Nzones)))
    r_ins = np.repeat(r_in_lst,int(Nc/Nzones))
    ns = np.repeat(n_lst,(int(Nc/Nzones)))
    for i in range(Nc):
        #print(r_ins[i],ns[i],model.parameters[model.model_options['cs'][i]].value,ps[i])
        surfdens_profile =  ns[i] * 10.0**model.parameters[model.model_options['cs'][i]].value*(radius/r_ins[i])**ps[i]
        surfdens_profiles.append(surfdens_profile)

    Nk = int(len(kappas)/Nzones)
    
    radius_cm = radius*14959787070000.0
    for i in range(Nzones):
        zone_idx = np.logical_and(radius >= r_in_lst[i],radius < r_out_lst[i]) 
        mass_in_zone = 0.0
        for j in range(i*Nk,(i+1)*Nk):
            mass_comp = np.trapz(2*np.pi*radius_cm[zone_idx]*surfdens_profiles[j][zone_idx], radius_cm[zone_idx])
            print('Dust mass in zone %d, comp %d = %.4e g = %.4e M_Earth = %.4e M_Sun'%(i+1,j+1,mass_comp,mass_comp/5.97219e27,mass_comp/1.9891e33))
            mass_in_zone += mass_comp
        print('Dust mass in zone %d = %.4e g = %.4e M_Earth = %.4e M_Sun'%(i+1,mass_in_zone,mass_in_zone/5.97219e27,mass_in_zone/1.9891e33))
    
def calc_fluxes(wave,radius,parameters):
    #with crystals: r_in,r_gap_in,r_gap_out,r_cryst,tprofile,surfdens_profile,kappa_inner,kappa_cryst,kappa_outer,dscale,rstarau,tstar,distance,cos_i,stellar_SED = parameters
    r_in_lst,r_out_lst,tprofile,surfdens_profiles,kappas,dscale,rstarau,tstar,distance,cos_i,inner_rim_boost_factor,stellar_SED = parameters
    #r_in_lst,r_out_lst,tprofile,surfdens_profiles,kappas,dscale,rstarau,tstar,distance,cos_i,f_outburst,stellar_SED = parameters
   
    #print('qi',q_inner)
    #print('qo',q_outer)
    Nzones = len(r_in_lst)
    flux =     np.zeros((Nzones,len(wave),len(radius)))
    fluxtot =  np.empty((Nzones,len(wave)))
    fluxstar = np.empty((len(wave)))
    surfdens = np.zeros((Nzones,len(radius)))
    tau =      np.zeros((Nzones,len(wave),len(radius)))
    
    freq_arr = np.transpose(np.broadcast_to((c0*1e6)/wave,(len(radius),len(wave))))
    #freq_arr = np.tile((c0*1e6)/wave,(np.shape(flux)[0],1))
    #radius_arr = np.transpose(np.tile(radius,(np.shape(flux)[1],1)))
    radius_arr = np.broadcast_to(radius,np.shape(flux)[1:])
    #tprofile_arr = np.transpose(np.tile(tprofile,(np.shape(flux)[1],1)))
    tprofile_arr = np.broadcast_to(tprofile,np.shape(flux)[1:])
    #print(flux.shape,freq_arr.shape,radius_arr.shape)

    Nk = int(len(kappas)/Nzones)
    for i in range(Nzones):
        zone_idx = np.logical_or(radius_arr < r_in_lst[i],radius_arr >= r_out_lst[i]) 
        tau_arr = 0.0*tprofile_arr
        surfdens_arr = 0.0*tprofile_arr
        surfdens_arr_int = 0.0*tprofile_arr
        # calculate the flux as a function of radius and wavelength
        for j in range(i*Nk,(i+1)*Nk):
            #print('in',i,np.max(surfdens_profiles[i]),np.max(kappas[i]))
            #surfdens_arr = np.transpose(np.tile(surfdens_profiles[j],(np.shape(flux)[1],1)))
            surfdens_arr = np.broadcast_to(surfdens_profiles[j],np.shape(flux)[1:]) 
            #kappa_arr = np.tile(kappas[j],(np.shape(flux)[0],1))
            kappa_arr = np.transpose(np.broadcast_to(kappas[j],(len(radius),len(wave))))
            tau_arr += surfdens_arr*kappa_arr
            surfdens_arr_int += surfdens_arr
            
        # write total kappa to file
        # kappa_sum = tau_arr/surfdens_arr_int
        # f = open('/allegro6/matisse/varga/targets3/HD_144432/modelling/kappa_new_z%d.txt'%(i+1),'w')
        # f.write('# lambda(um) kappa(cm^2 g^-1)\n')
        # for ii in range(len(wave)):
            # f.write('%.3f %.3f\n'%(wave[ii],kappa_sum[ii,0]))
        # f.close

        fluxtemp = (1.0-np.exp(-tau_arr))*(B_nu(tprofile_arr,freq_arr))
        #fluxtemp = tau_arr*(B_nu(tprofile_arr,freq_arr))
        fluxtemp[zone_idx] = 0.0
        surfdens_arr_int[zone_idx] = 0.0
        tau_arr[zone_idx] = 0.0
        flux[i,:,:] = fluxtemp
        surfdens[i,:] = surfdens_arr_int[0,:]
        tau[i,:,:] = tau_arr
     
    #flux[0,:,0] = inner_rim_boost_factor*flux[0,:,0]
    '''    
    fig, ax = plt.subplots(1, figsize=(8,6), sharex=True)
    for i in range(int(Nzones)):
        print(surfdens_arr.shape)
        zone_idx = np.logical_and(radius > r_in_lst[i],radius <= r_out_lst[i]) 
        surfdens_arr = 0.0*np.copy(tprofile_arr)
        for j in range(i*Nk,(i+1)*Nk):
            ax.plot(radius[zone_idx],surfdens_profiles[j][zone_idx],
                label=model.model_options['qs'][j-i*Nk] )
            
    surfdensint = np.sum(surfdens,axis=0)
    ax.plot(radius,surfdensint,'-k' )    
    ax.set_xlabel(r'$r$ (au)')
    ax.set_ylabel(r'$\Sigma\ (\mathrm{g}\,\mathrm{cm}^{-2})$')
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.legend(fontsize=6)
    plt.show()
    '''
    '''
    tau_arr = 0.0*tprofile_arr
    for i in range(int(Nk/2),Nk):
        #print('out',i,np.max(surfdens_profiles[i]),np.max(kappas[i]))
        surfdens_arr = np.transpose(np.tile(surfdens_profiles[i],(np.shape(flux)[1],1)))
        kappa_arr = np.tile(kappas[i],(np.shape(flux)[0],1))
        tau_arr += surfdens_arr*kappa_arr
  
    fluxtemp = (1.0-np.exp(-tau_arr))*(B_nu(tprofile_arr,freq_arr))  
    #fluxtemp[cryst_idx] = (1.0-np.exp(-surfdens_arr[cryst_idx]*kappa_cryst_arr[cryst_idx]))*(B_nu(tprofile_arr[cryst_idx],freq_arr[cryst_idx]))  
    fluxtemp[in_idx] = 0.0
    '''
    '''
    fluxtemp[gap_idx] = 0.0
    tsurface[gap_idx_1D] = 0.0
    
    in_idx = radius_arr < r_in
    edge_idx = find_nearest_idx(radius,r_in)+1
    flux_edge = np.tile(fluxtemp[edge_idx,:],(np.shape(flux)[0],1))  
    #print(fluxtemp.shape,fluxtemp[in_idx].shape,flux_edge.shape)
    fluxtemp[in_idx] = fluxtemp[in_idx] + flux_edge[in_idx]*np.exp(-( (radius_arr[in_idx]-r_in)**2 / ( 2.0 * (dr_smooth/(2.0*np.sqrt(2.0*np.log(2.0))))**2 ) ) )    

    gap_inner_edge_idx = find_nearest_idx(radius,r_gap_in)-1
    flux_edge = np.tile(fluxtemp[gap_inner_edge_idx,:],(np.shape(flux)[0],1))  
    fluxtemp[gap_idx] = fluxtemp[gap_idx] + flux_edge[gap_idx]*np.exp(-( (radius_arr[gap_idx]-r_gap_in)**2 / ( 2.0 * (dr_smooth/(2.0*np.sqrt(2.0*np.log(2.0))))**2 ) ) )    
    
    gap_outer_edge_idx = find_nearest_idx(radius,r_gap_out)+1
    #print(gap_outer_edge_idx)
    flux_edge = np.tile(fluxtemp[gap_outer_edge_idx,:],(np.shape(flux)[0],1))  
    #print(fluxtemp[gap_outer_edge_idx+1,:])
    fluxtemp[gap_idx] = fluxtemp[gap_idx] + flux_edge[gap_idx]*np.exp(-( (radius_arr[gap_idx]-r_gap_out)**2 / ( 2.0 * (dr_smooth/(2.0*np.sqrt(2.0*np.log(2.0))))**2 ) ) )    
    '''
    #flux[:,:,1] = fluxtemp
    radius_arr_arr = np.broadcast_to(radius_arr,np.shape(flux))
    #print(radius_arr.shape,radius_arr_arr.shape,flux.shape)

    # now integrate over the disk surface to get the total flux spectrum:
    fluxtot = np.trapz(2*np.pi*radius_arr_arr*flux, radius_arr_arr,axis=-1)
    #print(flux.shape,fluxtot.shape)

    '''
    for j in range(0,len(flux[0,:])):
        sumflmp = 0.0
        sumflsur = 0.0
        #sumflinnerrim = 0.0
        for i in range(0,len(radius)-1):
            twopr = 2.0*np.pi*radius[i]
            sumflmp = sumflmp + (flux[i,j,0]*(radius[i+1]-radius[i]))*twopr
            sumflsur = sumflsur + (flux[i,j,1]*(radius[i+1]-radius[i]))*twopr
            #sumflinnerrim = sumflinnerrim + (flux[i,j,2]*(radius[i+1]-radius[i]))*twopr
        fluxtot[j,0] = sumflmp
        fluxtot[j,1] = sumflsur
        #fluxtot[j,2] = sumflinnerrim
    '''

    # add stellar flux: 
    if stellar_SED is not None:
        fluxstar = stellar_SED/1e26*dscale**2 #*f_outburst
    else:
        #black body assumed here, stellar model used in MCMAX is Kurucz!!
        fluxstar = rstarau**2*np.pi*B_nu(tstar,(c0*1e6)/wave) #*f_outburst #0.0*wave # 
        
    # sum up all the intensities for the zones:
    fluxint = np.sum(flux,axis=0)
    tauint = np.sum(tau,axis=0)
    surfdensint = np.sum(surfdens,axis=0)
    total_flux_arr_Jy = 1e26*fluxtot*cos_i/dscale**2 #before 2023.07.13
    #total_flux_arr_Jy = 1e26*fluxtot/dscale**2 #after 2023.07.13
    #fluxint = (flux[:,:,0] + flux[:,:,1]) # + flux[:,:,2])
    #total_flux0_Jy = 1e26*fluxtot[:,0]*cos_i/dscale**2
    #total_flux1_Jy = 1e26*fluxtot[:,1]*cos_i/dscale**2
    #total_flux2_Jy = 1e26*fluxtot[:,2]*cos_i/dscale**2
    fluxstar_Jy = 1e26*(fluxstar)/dscale**2
    total_flux_Jy = fluxstar_Jy + np.sum(total_flux_arr_Jy,axis=0)
    #total_flux_Jy = fluxstar_Jy + total_flux0_Jy + total_flux1_Jy #+ total_flux2_Jy #check validity of cos_i!!!!
    
    # print(radius_arr)
    rr = radius_arr/distance/3600.0*np.pi/180.0
    ll = np.transpose(np.broadcast_to(wave,(len(radius),len(wave))))
    
    #return fluxint,rr,ll,total_flux_Jy,fluxstar_Jy,total_flux0_Jy, total_flux1_Jy, 0.0*total_flux1_Jy,0.0*tprofile,tprofile
    return fluxint,rr,ll,total_flux_Jy,fluxstar_Jy,total_flux_arr_Jy, tprofile, tauint, surfdensint


def image_TGMdust(model):
    distance = model.parameters['dist_pc'].value # in pc
    tstar = model.parameters['T_star'].value # in K
    lstar = model.parameters['L_star'].value # in solar luminosity
    rstarau = model.parameters['R_star'].value*rsun/au
    dscale = distance*parsec/au
    if model.model_options['use_powerlaw_tempgrad']:
        T_in = model.parameters['T_i1'].value
    #f_outburst = model.parameters['f_outburst'].value
    inner_rim_boost_factor = 1.0 #model.parameters['inner_rim_boost_factor'].value
    
    # disk parameters
    rtaper= model.parameters['r_taper_au'].value # in AU
    #t_subl = model.parameters['T_sub'].value # maximum dust temperature in K
    Nzones = model.model_options['Nzones']
    r_in_lst = []
    r_out_lst = []
    p_lst = []
    n_lst = []
    for i in range(Nzones):
        r_in_lst.append(model.parameters['r_i%d'%(i+1)].value)
        p_lst.append(model.parameters['p%d'%(i+1)].value)
        n_lst.append(10.0**model.parameters['n%d'%(i+1)].value)
        if i>0:
            r_out_lst.append(r_in_lst[-1])
    #print(n_lst)
    r_out_lst.append(rtaper)   
    r_in = r_in_lst[0]
    cos_i = model.parameters['axis_ratio'].value
    #C_i = model.parameters['C_i'].value
    #C_o = model.parameters['C_o'].value
    if 'kappa0' in model.parameters:
        kappa0 = model.parameters['kappa0'].value
    Nc = len(model.model_options['cs'])
 
    if 'kappa0' in model.parameters:
        pki = model.model_options['parkappa_idx']
        for i in range(Nzones):
            lambda_break = model.parameters['lambda_break_z%d'%(i+1)].value
            p0_kappa = model.parameters['p0_kappa_z%d'%(i+1)].value
            p1_kappa = model.parameters['p1_kappa_z%d'%(i+1)].value
            f = models.SmoothlyBrokenPowerLaw1D(amplitude=kappa0, x_break=lambda_break,
                                        alpha_1=p0_kappa, alpha_2=p1_kappa,delta=0.2)
            
            #print(model.parameters['sdens_z%d_%d'%((i+1),pki)].wl_sed_q)
            #print(f(model.parameters['sdens_z%d_%d'%((i+1),pki)].wl_sed_q))
            model.parameters['sdens_z%d_%d'%((i+1),pki)].q = f(model.parameters['sdens_z%d_%d'%((i+1),pki)].wl_q) 
            model.parameters['sdens_z%d_%d'%((i+1),pki)].sed_q = f(model.parameters['sdens_z%d_%d'%((i+1),pki)].wl_sed_q)
            model.model_options['temp_cal_pm'][pki+i*int(Nc/Nzones)]['opac_orig'] = f(model.model_options['temp_cal_pm'][pki+i*int(Nc/Nzones)]['wl_orig'])
            idx = model.model_options['cs'][pki+i*int(Nc/Nzones)]
            model.parameters[idx].q = f(model.parameters[idx].wl_q)
            model.parameters[idx].sed_q = f(model.parameters[idx].wl_sed_q)
            
    kappas = []
    kappas_sed = []
    for i in range(Nc):
        kappas.append(model.parameters[model.model_options['cs'][i]].q)
        if model.model_options['fit_SED']:
            kappas_sed.append(model.parameters[model.model_options['cs'][i]].sed_q)
        #print(i,model.model_options['qs'][i],np.max(kappas[i]),np.max(kappas_sed[i]))

    '''
    if show_plots:
        print(model.model_options['qs'])
        fig, axes = plt.subplots(int(Nc/Nzones), figsize=(6,Nc/Nzones*5), sharex=True)
        for i in range(int(Nc/Nzones)):
            if Nc > 1:
                ax = axes[i]
            else:
                ax = axes
            #plt.plot(model.parameters[c_i[0]].wl_sed_q,model.parameters[c_i[0]].sed_q,marker='.')
            for j in range(Nzones):
                ij = j*int(Nc/Nzones)+i
                ax.plot(model.parameters[model.model_options['cs'][ij]].wl_q,model.parameters[model.model_options['cs'][ij]].q,
                    marker='x',label=model.model_options['qs'][ij] )
            #plt.plot(model.parameters[c_i[0]].wl_sed_q,model.model_options['stellar_flux_multiwl']*1e2)
            ax.set_xlabel(r'$\lambda$ ($\mu$m)')
            ax.set_ylabel(r'$\kappa\ (\mathrm{cm}^2\,\mathrm{g}^{-1})$')
            ax.legend()
        plt.show()
    '''
        
    # calculate the innermost dust radius:
    #rinner = 0.06867894 * (np.sqrt(lstar))*((1500/t_subl)**2)
    rinner = r_in
    '''
    rinner = r_in-dr_smooth*2.0
    if rinner < 0.01:
        rinner = 0.01
    '''
    # fill the radius grid
    radius = np.logspace(np.log10(rinner),np.log10(rtaper),model.model_options['nx'], endpoint='true')

    # fill the temperature grid
    #tmidplane = radius * 0.0
    #tmidplane = t_subl*(radius/rinner)**tg_mp_out
    
    #derive the temperature profile
    Tdust_init = 1500.0
    if not model.model_options['use_powerlaw_tempgrad']:
        func = lambda Tdust,rd,kappa,leftside,nu : (1/rd)**2 * leftside - \
            4.0 * np.trapz(kappa*B_nu(Tdust,nu),x=nu)
        tprofile = []
        #leftside_in = model.model_options['leftside_in']
        #tprofile.append(fsolve(func, Tdust_init,(radius[0],model.model_options['opac_orig_in'],model.model_options['leftside_in'],model.model_options['nu_orig'])))
        #leftside_in = (model.parameters['R_star'].value*0.00465047)**2*np.trapz(f(model.model_options['wl_orig'])*model.model_options['Bnu_star'],x=model.model_options['nu_orig'])
        #tprofile.append(fsolve(func, Tdust_init,(radius[0],f(model.model_options['wl_orig']),leftside_in,model.model_options['nu_orig'])))
        i=0
        #print(Nzones,r_out_lst)
        for kk in range(Nzones):
            avg_opac_orig = 0.0*model.model_options['temp_cal_pm'][0]['opac_orig']
            for j in range(kk*int(Nc/Nzones),(kk+1)*int(Nc/Nzones)):
                avg_opac_orig += model.model_options['temp_cal_pm'][j]['opac_orig']*10.0**model.parameters[model.model_options['cs'][j]].value
            leftside = (model.parameters['R_star'].value*0.00465047)**2*np.trapz(avg_opac_orig*
                model.model_options['temp_cal_pm'][0]['Bnu_star'],x=model.model_options['temp_cal_pm'][0]['nu_orig'])
            tprofile.append(fsolve(func, Tdust_init,(radius[i],avg_opac_orig,leftside,
                model.model_options['temp_cal_pm'][0]['nu_orig'])))
            i+=1
            while i<(len(radius)-1) and radius[i] < r_out_lst[kk]: #r_cryst:
                #print(kk,i,len(radius))
                #print(radius[i],r_out_lst[kk])
                Tdust_init = tprofile[-1]
                tprofile.append(fsolve(func, Tdust_init,(radius[i],avg_opac_orig,leftside,
                    model.model_options['temp_cal_pm'][0]['nu_orig'])))
                i+=1
        Tdust_init = tprofile[-1]
        tprofile.append(fsolve(func, Tdust_init,(radius[i],avg_opac_orig,leftside,
            model.model_options['temp_cal_pm'][0]['nu_orig'])))
        #print(i,len(radius),len(tprofile))
        
        '''
        avg_opac_orig_in = 0.0*model.model_options['temp_cal_pm'][0]['opac_orig']
        avg_opac_orig_out = 0.0*model.model_options['temp_cal_pm'][2]['opac_orig']
        for i in range(int(Nc/2)):
            avg_opac_orig_in += model.model_options['temp_cal_pm'][i]['opac_orig']*model.parameters[model.model_options['cs'][i]].value
        for i in range(int(Nc/2),Nc):
            avg_opac_orig_out += model.model_options['temp_cal_pm'][i]['opac_orig']*model.parameters[model.model_options['cs'][i]].value

        leftside_in = (model.parameters['R_star'].value*0.00465047)**2*np.trapz(avg_opac_orig_in*
            model.model_options['temp_cal_pm'][0]['Bnu_star'],x=model.model_options['temp_cal_pm'][0]['nu_orig'])
        leftside_out = (model.parameters['R_star'].value*0.00465047)**2*np.trapz(avg_opac_orig_out*
            model.model_options['temp_cal_pm'][2]['Bnu_star'],x=model.model_options['temp_cal_pm'][2]['nu_orig'])
        
        tprofile.append(fsolve(func, Tdust_init,(radius[0],avg_opac_orig_in,leftside_in,
            model.model_options['temp_cal_pm'][0]['nu_orig'])))
        i=1
        while radius[i] < r_gap_out: #r_cryst:
            Tdust_init = tprofile[-1]
            #tprofile.append(fsolve(func, Tdust_init,(radius[i],model.model_options['opac_orig_in'],model.model_options['leftside_in'],model.model_options['nu_orig'])))
            #tprofile.append(fsolve(func, Tdust_init,(radius[i],f(model.model_options['wl_orig']),leftside_in,model.model_options['nu_orig'])))
            tprofile.append(fsolve(func, Tdust_init,(radius[i],avg_opac_orig_in,leftside_in,
                model.model_options['temp_cal_pm'][0]['nu_orig'])))
            i+=1

        j = i
        for i in range(j,len(radius)): 
            Tdust_init = tprofile[-1] 
            tprofile.append(fsolve(func, Tdust_init,(radius[i],avg_opac_orig_out,leftside_out,
                model.model_options['temp_cal_pm'][2]['nu_orig'])))
            #tprofile.append(fsolve(func, Tdust_init,(radius[i],avg_opac_orig_out,model.model_options['temp_cal_pm'][3]['leftside'],model.model_options['temp_cal_pm'][3]['nu_orig'])))
        '''
        '''
        plt.figure()
        plt.plot(radius,tprofile,marker='.')
        plt.yscale('log')
        plt.xscale('log')
        plt.show()
        '''

        tprofile = np.squeeze(np.array(tprofile))
    else:
        tprofile = np.array(T_in*(radius/radius[0])**model.parameters['tg_sur'].value)
        
    surfdens_profiles = []
    ps = np.repeat(p_lst,(int(Nc/Nzones)))
    r_ins = np.repeat(r_in_lst,int(Nc/Nzones))
    ns = np.repeat(n_lst,(int(Nc/Nzones)))
    for i in range(Nc):
        #print(r_ins[i],ns[i],model.parameters[model.model_options['cs'][i]].value,ps[i])
        surfdens_profile =  ns[i] * 10.0**model.parameters[model.model_options['cs'][i]].value*(radius/r_ins[i])**ps[i]
        surfdens_profiles.append(surfdens_profile)
        #print(i,np.nanmax(surfdens_profile))

    #surfdens_profile = model.parameters['c_dustkappa_i'].value*(radius/r_in)**p_in
    #out_idx = radius > r_gap_out
    #surfdens_profile[out_idx] = model.parameters['c_dustkappa_o'].value*(radius[out_idx]/r_gap_out)**p_out
    #cryst_idx = np.logical_and(radius > r_cryst, radius < r_gap_out) 
    #surfdens_profile[cryst_idx] = model.parameters['c_dustkappa_c'].value*(radius[cryst_idx]/r_cryst)**p_cryst
    
    #gap_idx = np.logical_and(radius > r_gap_in, radius < r_gap_out) 
    #surfdens_profile[gap_idx] = 0.0
    #rdust = 0.06867894 * (np.sqrt(lstar))
    #tsurface = 1500.0*(radius/rdust)**tg_sur_out
    #tsurface[tsurface>1500.0]=0.0
    #tsurface = t_in*(radius/r_in)**tg_sur_in
    #out_idx = radius > r_gap_out
    #tsurface[out_idx] = t_out*(radius[out_idx]/r_gap_out)**tg_sur_out
    #in_idx = radius < r_in
    #tsurface[in_idx] = 0.0
    
    #q_inner = q_inner*C_i
    #q_outer = q_outer*C_o
    #if model.fit_SED:
    #    q_sed_outer = q_sed_outer*C_o
        
    if 'stellar_flux_interf' in model.model_options:
        stellar_SED = model.model_options['stellar_flux_interf']
    else:
        stellar_SED = None    
    
    #parameters = [r_gap_in,r_gap_out,tmidplane,tsurface,q_inner,q_outer,qjump_width_factor,r_qjump, qjump_width_factor,dscale,rstarau,tstar,distance,cos_i,tau_in,stellar_SED,r_gap_mp_out ]
    #parameters = [r_in,r_gap_in,r_gap_out,r_cryst,tprofile,surfdens_profile,kappa_inner,kappa_cryst,kappa_outer,dscale,rstarau,tstar,distance,cos_i,stellar_SED]
    parameters = [r_in_lst,r_out_lst,tprofile,surfdens_profiles,kappas,dscale,rstarau,tstar,distance,cos_i,inner_rim_boost_factor,stellar_SED ]
    #parameters = [r_in_lst,r_out_lst,tprofile,surfdens_profiles,kappas,dscale,rstarau,tstar,distance,cos_i,f_outburst,stellar_SED ]
    wave = model.parameters[model.model_options['cs'][0]].wl_q #um
    #fluxint,rr,ll,total_flux_Jy,fluxstar_Jy,total_flux0_Jy, total_flux1_Jy, total_flux2_Jy, tmidplane, tsurface = calc_fluxes(wave,radius,parameters)
    fluxint,rr,ll,total_flux_Jy,fluxstar_Jy,total_flux_arr_Jy, tprofile_v, tauint, surfdensint = calc_fluxes(wave,radius,parameters)
    
    if model.fit_SED:
        if 'stellar_flux_multiwl' in model.model_options:
            stellar_SED = model.model_options['stellar_flux_multiwl']
        else:
            stellar_SED = None
        parameters[4] = kappas_sed
        #parameters[7] = kappa_sed_cryst
        parameters[-1] = stellar_SED
        
        wave_sed = model.parameters[model.model_options['cs'][0]].wl_sed_q #um
        #print('wl',wave_sed)
        #fluxint_sed,rr_sed,ll_sed,total_flux_Jy_sed,fluxstar_Jy_sed,total_flux0_Jy, total_flux1_Jy, total_flux2_Jy, tmidplane,tsurface = calc_fluxes(wave_sed,radius,parameters,model)
        fluxint_sed,rr_sed,ll_sed,total_flux_Jy_sed,fluxstar_Jy_sed,total_flux_arr_Jy, tprofile_v, tauint_sed, surfdensint_sed = calc_fluxes(wave_sed,radius,parameters)
    '''
    if show_plots:
        plt.loglog(wave,1e26*(fluxstar + fluxtot[:,0] + fluxtot[:,1]+ fluxtot[:,2])/dscale**2,label='total')
        plt.loglog(wave,1e26*(fluxtot[:,0]/dscale**2),label='midplane')
        plt.loglog(wave,1e26*(fluxtot[:,1]/dscale**2),label='surface')
        plt.loglog(wave,1e26*(fluxtot[:,2]/dscale**2))
        plt.loglog(wave,1e26*fluxstar/dscale**2,label='star')
        plt.xlim([0.3,400])
        plt.ylim([1e-2,50])
        plt.xlabel('wavelength ($\mu$m)')
        plt.ylabel('Flux (Jy)')
        plt.legend(loc='best')
        plt.show()

        plt.loglog(radius,fluxint[:,40]*radius**2,label='$\lambda$ = ' + str('%6.2f' % wave[40]))
        plt.loglog(radius,fluxint[:,150]*radius**2,label='$\lambda$ = ' + str('%6.2f' %wave[150]))
        plt.loglog(radius,fluxint[:,250]*radius**2,label='$\lambda$ = ' + str('%6.2f' %wave[250]))
        plt.loglog(radius,fluxint[:,350]*radius**2,label='$\lambda$ = ' + str('%6.2f' %wave[350]))
        plt.loglog(radius,fluxint[:,450]*radius**2,label='$\lambda$ = ' + str('%6.2f' %wave[450]))
        plt.legend(loc='best')
        plt.xlabel('radius (AU)')
        plt.ylabel('Intensity*r$^{2}$')
        plt.ylim((1e-14,1e-6))
        plt.show()
    '''

    if model.fit_SED == False:
        #return fluxint,rr,ll,total_flux_Jy,fluxstar_Jy,total_flux0_Jy, total_flux1_Jy, total_flux2_Jy, tmidplane, tsurface
        return fluxint,rr,ll,total_flux_Jy,fluxstar_Jy,total_flux_arr_Jy, tprofile, tauint, surfdensint
    else: 
        #return fluxint,rr,ll,total_flux_Jy,fluxstar_Jy,total_flux_Jy_sed,fluxstar_Jy_sed,total_flux0_Jy, total_flux1_Jy, total_flux2_Jy, tmidplane, tsurface
        return fluxint,rr,ll,total_flux_Jy,fluxstar_Jy,total_flux_Jy_sed,fluxstar_Jy_sed,total_flux_arr_Jy, tprofile, tauint, tauint_sed, surfdensint, surfdensint_sed, fluxint_sed, rr_sed, ll_sed
#write to an ascii output file:

# f = open('intensity.txt','w')

#x=arange(0,2*pi,0.01)

#In [13]: for val in x:
#   ....:     f.write("{:.4f} {:7.4f}\n".format(radius[i], fluxint[i,7]))

# f.write("{:.4f} {:7.4f} {:7.4f} {:7.4f}\n".format(wave[7],wave[43],wave[58],wave[70]))
# for i in range(0,len(radius)):
#     f.write("{:.4f} {:7.4e} {:7.4e} {:7.4e} {:7.4e}\n".format(radius[i], fluxint[i,7],fluxint[i,43],fluxint[i,58],fluxint[i,70]))

# f.close()

def find_nearest_idx(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx