# Copyright (C) 2025 Xiaoyang Chen - All Rights Reserved
# Licensed under the GNU GENERAL PUBLIC LICENSE Version 3
# Contact: xiaoyang.chen.cz@gmail.com

import numpy as np
from copy import deepcopy as copy
from astropy.io import fits
import astropy.units as u
import astropy.constants as const
from astropy.cosmology import WMAP9 as cosmo
from scipy.interpolate import RegularGridInterpolator

class TorusFrame(object): 
    # SKIRTor, https://sites.google.com/site/skirtorus/sed-library
    def __init__(self, filename=None, w_min=None, w_max=None, 
                 cframe=None, v0_redshift=None, flux_scale=None): # , spec_R_inst=None
        
        self.filename = filename        
        self.w_min = w_min # currently not used
        self.w_max = w_max # currently not used
        self.cframe = cframe 
        self.v0_redshift = v0_redshift
        self.flux_scale = flux_scale
        # self.spec_R_inst = spec_R_inst
        
        self.num_coeffs = 1 # only have one norm-free model in a single fitting component (since disc and torus are tied)
        self.read_skirtor()
        
    def read_skirtor(self): 
        # skirtor_disc = np.loadtxt(self.file_disc) # [n_wave_ini+6, n_tau*n_oa*n_rrat*n_incl+1]
        # skirtor_torus = np.loadtxt(self.file_dust) # [n_wave_ini+6, n_tau*n_oa*n_rrat*n_incl+1]
        skirtor_disc = fits.open(self.filename)[0].data[0]
        skirtor_torus = fits.open(self.filename)[0].data[1]

        wave = skirtor_disc[6:-1,0] # 1e-3 to 1e3 um; omit the last one with zero-value SED
        n_tau = 5; tau = np.array([ 3, 5, 7, 9, 11 ])
        n_oa = 8; oa = np.array([ 10, 20, 30, 40, 50, 60, 70, 80 ])
        n_rrat = 3; rrat = np.array([ 10, 20, 30 ])
        n_incl = 10; incl = np.array([ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90 ])
                
        disc  = np.zeros([n_tau, n_oa, n_rrat, n_incl, len(wave)]) 
        torus = np.zeros([n_tau, n_oa, n_rrat, n_incl, len(wave)]) 
        mass  = np.zeros([n_tau, n_oa, n_rrat]) # torus dust mass
        eb    = np.zeros([n_tau, n_oa, n_rrat]) 
        # All spectra are given in erg/s/um, normalized to disc lum of 1 Lsun
        # Not that the spectra in lum unit should be considered as flux * 4pi * dl2, 
        # where flux is depended on viewing angle and 
        # the 1 Lsun normalization is integrated with anisotropic flux function.
        # Dust mass in Msun
        # eb is energy balance ratio of torus, i.e., inclination integrated Lum_torus/Lum_AGN(intrinsic)

        for i_tau in range(n_tau):
            for i_oa in range(n_oa):
                for i_rrat in range(n_rrat):
                    for i_incl in range(n_incl):
                        mask = skirtor_disc[0,:] == tau[i_tau] 
                        mask &= skirtor_disc[1,:] == oa[i_oa] 
                        mask &= skirtor_disc[2,:] == rrat[i_rrat] 
                        mask &= skirtor_disc[3,:] == incl[i_incl] 
                        mass[i_tau, i_oa, i_rrat] = skirtor_torus[4,mask][0]
                        eb[i_tau, i_oa, i_rrat] = skirtor_torus[5,mask][0]
                        disc[i_tau, i_oa, i_rrat, i_incl, :] = skirtor_disc[6:-1,mask][:,0]
                        torus[i_tau, i_oa, i_rrat, i_incl, :] = skirtor_torus[6:-1,mask][:,0]
                        # in the original library the torus sed and mass is normalized to Lum_AGN of 1 Lsun, 
                        # here renormlized them to Lum_Torus of 1e12 Lsun (i.e., Lum_AGN = 1e12 Lsun / EB_Torus) 
                        disc[i_tau, i_oa, i_rrat, i_incl, :]  *= 1e12 / eb[i_tau, i_oa, i_rrat]
                        torus[i_tau, i_oa, i_rrat, i_incl, :] *= 1e12 / eb[i_tau, i_oa, i_rrat]
                        mass[i_tau, i_oa, i_rrat] *= 1e12 / eb[i_tau, i_oa, i_rrat]
        
        # convert unit: 1 erg/s/um -> flux_scale * erg/s/AA/cm2
        wave *= 1e4
        lum_dist = cosmo.luminosity_distance(self.v0_redshift).to('cm').value
        lum_area = 4*np.pi * lum_dist**2 # in cm2
        disc *= 1e-4 / lum_area / self.flux_scale
        torus *= 1e-4 / lum_area / self.flux_scale
        disc[disc <= 0]   = disc[disc>0].min()
        torus[torus <= 0] = torus[torus>0].min()
        
        # for interpolation
        ini_pars = (tau, oa, rrat, incl, np.log10(wave))    
        fun_logdisc = RegularGridInterpolator(ini_pars, np.log10(disc), method='linear', bounds_error=False)
        fun_logtorus = RegularGridInterpolator(ini_pars, np.log10(torus), method='linear', bounds_error=False)
        # set bounds_error=False to avoid error by slight exceeding of x-val generated by least_square func
        # but do not use pars outside of initial range
        ini_pars = (tau, oa, rrat)    
        fun_mass = RegularGridInterpolator(ini_pars, mass, method='linear', bounds_error=False)
        fun_eb = RegularGridInterpolator(ini_pars, eb, method='linear', bounds_error=False)

        self.skirtor = {'tau':tau, 'oa':oa, 'rratio':rrat, 'incl':incl, 
                        'wave':wave, 'logwave':np.log10(wave), 
                        'disc':disc, 'fun_logdisc':fun_logdisc, 
                        'torus':torus, 'fun_logtorus':fun_logtorus, 
                        'mass':mass, 'fun_mass':fun_mass, 
                        'eb':eb, 'fun_eb':fun_eb } 
        
    def get_info(self, tau, oa, rratio):
        fun_mass = self.skirtor['fun_mass']
        fun_eb = self.skirtor['fun_eb']
        gen_pars = np.array([tau, oa, rratio])
        gen_mass = fun_mass(gen_pars)
        gen_eb   = fun_eb(gen_pars)
        return gen_mass, gen_eb

    def models_unitnorm_obsframe(self, wavelength, input_pars, if_pars_flat=True):
        # pars: voff (to adjust redshift), tau, oa, rratio, incl
        # comps: 'disc', 'torus', 'disc+torus'
        if if_pars_flat: 
            pars = self.cframe.flat_to_arr(input_pars)
        else:
            pars = copy(input_pars)

        for i_comp in range(pars.shape[0]):
            tau, oa, rratio, incl = pars[i_comp,:][1:5]
            
            # interpolate model for given pars in initial wavelength (rest)
            ini_logwave = self.skirtor['logwave'].copy()
            fun_logdisc = self.skirtor['fun_logdisc']
            fun_logtorus = self.skirtor['fun_logtorus']
            gen_pars = np.array([[tau, oa, rratio, incl, w] for w in ini_logwave]) # gen: generated
            if np.isin('disc', self.cframe.info_c[i_comp]['mod_used']):
                gen_logdisc = fun_logdisc(gen_pars)
            if np.isin('dust', self.cframe.info_c[i_comp]['mod_used']):
                gen_logtorus = fun_logtorus(gen_pars)    

            # redshifted to obs-frame
            ret_logwave = np.log10(wavelength) # in AA
            z_ratio = (1 + self.v0_redshift) * (1 + pars[i_comp,0]/299792.458) # (1+z) = (1+zv0) * (1+v/c)            
            ini_logwave += np.log10(z_ratio)
            if np.isin('disc', self.cframe.info_c[i_comp]['mod_used']):
                gen_logdisc -= np.log10(z_ratio)
                ret_logdisc  = np.interp(ret_logwave, ini_logwave, gen_logdisc, 
                                         left=np.minimum(gen_logdisc.min(),-100), right=np.minimum(gen_logdisc.min(),-100))
            if np.isin('dust', self.cframe.info_c[i_comp]['mod_used']):
                gen_logtorus -= np.log10(z_ratio)
                ret_logtorus = np.interp(ret_logwave, ini_logwave, gen_logtorus, 
                                         left=np.minimum(gen_logtorus.min(),-100), right=np.minimum(gen_logtorus.min(),-100))
                
            # extended to longer wavelength
            mask_w = ret_logwave > ini_logwave[-1]
            if np.sum(mask_w) > 0:
                if np.isin('disc', self.cframe.info_c[i_comp]['mod_used']):
                    index = (gen_logdisc[-2]-gen_logdisc[-1]) / (ini_logwave[-2]-ini_logwave[-1])
                    ret_logdisc[mask_w] = gen_logdisc[-1] + index * (ret_logwave[mask_w]-ini_logwave[-1])
                if np.isin('dust', self.cframe.info_c[i_comp]['mod_used']):
                    index = (gen_logtorus[-2]-gen_logtorus[-1]) / (ini_logwave[-2]-ini_logwave[-1])
                    ret_logtorus[mask_w] = gen_logtorus[-1] + index * (ret_logwave[mask_w]-ini_logwave[-1])
                    
            if np.isin('disc', self.cframe.info_c[i_comp]['mod_used']):
                ret_disc = 10.0**ret_logdisc
                ret_disc[ret_logdisc <= -100] = 0
            if np.isin('dust', self.cframe.info_c[i_comp]['mod_used']):
                ret_torus = 10.0**ret_logtorus
                ret_torus[ret_logtorus <= -100] = 0
                
            models_scomp_obsframe = np.zeros_like(ret_logwave)
            if np.isin('disc', self.cframe.info_c[i_comp]['mod_used']): models_scomp_obsframe += ret_disc
            if np.isin('dust', self.cframe.info_c[i_comp]['mod_used']): models_scomp_obsframe += ret_torus
                
            models_scomp_obsframe = np.vstack((models_scomp_obsframe))
            models_scomp_obsframe = models_scomp_obsframe.T # add .T for a uniform format with other models with n_coeffs > 1
            
            if i_comp == 0: 
                models_mcomp_obsframe = models_scomp_obsframe
            else:
                models_mcomp_obsframe += models_scomp_obsframe # not tested
        return np.array(models_mcomp_obsframe) 

    ##########################################################################
    ########################### Output functions #############################

    def output_results(self, ff=None):
        num_mock_loops = ff.num_mock_loops
        best_chi_sq_l = ff.best_chi_sq
        best_x_lp = ff.best_fits_x
        best_coeff_lm = ff.best_coeffs
        fx0, fx1, fc0, fc1 = ff.model_index('torus', ff.full_model_type)

        num_torus_comps = self.cframe.num_comps
        num_torus_pars = self.cframe.num_pars
        num_torus_coeffs = self.num_coeffs

        output_torus_lcp = np.zeros((num_mock_loops, num_torus_comps, 1 + num_torus_pars + num_torus_coeffs ))
        # p: chi_sq, output_values, ssp_coeffs
        output_torus_lcp[:, :, 0]                  = best_chi_sq_l[:, None]
        output_torus_lcp[:, :, 1:(1+num_torus_pars)] = best_x_lp[:, fx0:fx1].reshape(num_mock_loops, num_torus_comps, num_torus_pars)
        output_torus_lcp[:, :, -num_torus_coeffs:]   = best_coeff_lm[:, fc0:fc1].reshape(num_mock_loops, num_torus_comps, num_torus_coeffs)
        self.output_torus_lcp = output_torus_lcp # save to model frame
        
        # output to screen
        output_torus_vals = {
            'mean': np.average(output_torus_lcp, weights=1/best_chi_sq_l, axis=0), 
            'rms' : np.std(output_torus_lcp, axis=0, ddof=1) }
        self.output_torus_vals = output_torus_vals # to print results
        self.print_results()
        
    def print_results(self):
        num_torus_comps = self.cframe.num_comps

        print('')
        for i_comp in range(num_torus_comps):
            print(f'Best-fit properties of torus component: <{self.cframe.comp_c[i_comp]}>')
            # msg  = f'| Voff (km/s)           = {self.output_torus_vals["mean"][1]:10.4f}'
            # msg += f' +/- {self.output_torus_vals["rms"][1]:0.4f}\n'
            msg  = f'| Opacity               = {self.output_torus_vals["mean"][i_comp,2]:10.4f}'
            msg += f' +/- {self.output_torus_vals["rms"][i_comp,2]:0.4f}\n'
            msg += f'| Out/in radii ratio    = {self.output_torus_vals["mean"][i_comp,4]:10.4f}'
            msg += f' +/- {self.output_torus_vals["rms"][i_comp,4]:0.4f}\n'
            msg += f'| Half OpenAng (degree) = {self.output_torus_vals["mean"][i_comp,3]:10.4f}'
            msg += f' +/- {self.output_torus_vals["rms"][i_comp,3]:0.4f}\n'
            msg += f'| Inclination (degree)  = {self.output_torus_vals["mean"][i_comp,5]:10.4f}'
            msg += f' +/- {self.output_torus_vals["rms"][i_comp,5]:0.4f}\n'
            msg += f'| Torus Lum (1e12 Lsun) = {self.output_torus_vals["mean"][i_comp,6]:10.4f}'
            msg += f' +/- {self.output_torus_vals["rms"][i_comp,6]:0.4f}'
            bar = '='*50
            print(bar)
            print(msg)
            print(bar)
