"""Molecular database (MDB) class

   * MdbExomol is the MDB for ExoMol
   * MdbHit is the MDB for HITRAN or HITEMP

"""
import numpy as np
import jax.numpy as jnp
import pathlib
from exojax.spec import hapi, exomolapi, exomol, atomllapi, atomll, hitranapi
from exojax.spec.hitran import gamma_natural as gn
import vaex

__all__ = ['MdbExomol','MdbHit','AdbVald','AdbKurucz']

class MdbExomol(object):
    """ molecular database of ExoMol

    MdbExomol is a class for ExoMol.

    Attributes:
        nurange: nu range [min,max] (cm-1)
        nu_lines (nd array): line center (cm-1)
        Sij0 (nd array): line strength at T=Tref (cm)
        dev_nu_lines (jnp array): line center in device (cm-1)
        logsij0 (jnp array): log line strength at T=Tref
        A (jnp array): Einstein A coeeficient
        gamma_natural (jnp array): gamma factor of the natural broadening
        elower (jnp array): the lower state energy (cm-1)
        gpp (jnp array): statistical weight
        jlower (jnp array): J_lower
        jupper (jnp array): J_upper
        n_Texp (jnp array): temperature exponent
        alpha_ref (jnp array): alpha_ref (gamma0)
        n_Texp_def: default temperature exponent in .def file, used for jlower not given in .broad
        alpha_ref_def: default alpha_ref (gamma0) in .def file, used for jlower not given in .broad

    """
    def __init__(self,path,nurange=[-np.inf,np.inf],margin=0.0,crit=-np.inf, bkgdatm="H2", broadf=True):
        """Molecular database for Exomol form

        Args:
           path: path for Exomol data directory/tag. For instance, "/home/CO/12C-16O/Li2015"
           nurange: wavenumber range list (cm-1) [min,max] or wavenumber grid
           margin: margin for nurange (cm-1)
           crit: line strength lower limit for extraction
           bkgdatm: background atmosphere for broadening. e.g. H2, He,
           broadf: if False, the default broadening parameters in .def file is used

        Note:
           The trans/states files can be very large. For the first time to read it, we convert it to HDF/vaex. After the second-time, we use the HDF5 format with vaex instead.

        """
        explanation_states="Note: Couldn't find the hdf5 format. We convert data to the hdf5 format. After the second time, it will become much faster."
        explanation_trans="Note: Couldn't find the hdf5 format. We convert data to the hdf5 format. After the second time, it will become much faster."

        self.path = pathlib.Path(path)
        t0=self.path.parents[0].stem
        molec=t0+"__"+str(self.path.stem)
        self.bkgdatm=bkgdatm
        print("Background atmosphere: ",self.bkgdatm)
        molecbroad=t0+"__"+self.bkgdatm

        self.crit = crit
        self.margin = margin
        self.nurange=[np.min(nurange),np.max(nurange)]
        self.broadf=broadf
        self.states_file = self.path/pathlib.Path(molec+".states.bz2")
        self.pf_file = self.path/pathlib.Path(molec+".pf")
        self.def_file = self.path/pathlib.Path(molec+".def")
        self.broad_file = self.path/pathlib.Path(molecbroad+".broad")

        if not self.def_file.exists():
                self.download(molec,extension=[".def"])
        if not self.pf_file.exists():
                self.download(molec,extension=[".pf"])
        if not self.states_file.exists():
                self.download(molec,extension=[".states.bz2"])
        if not self.broad_file.exists():
                self.download(molec,extension=[".broad"])

        #load def
        self.n_Texp_def, self.alpha_ref_def, self.molmass, numinf, numtag=exomolapi.read_def(self.def_file)

        #  default n_Texp value if not given
        if self.n_Texp_def is None:
            self.n_Texp_def=0.5
        #  default alpha_ref value if not given
        if self.alpha_ref_def is None:
            self.alpha_ref_def=0.07

        #load states
        if self.states_file.with_suffix(".bz2.hdf5").exists():
            states=vaex.open(self.states_file.with_suffix(".bz2.hdf5"))
            ndstates=vaex.array_types.to_numpy(states)
        else:
            print(explanation_states)
            states=exomolapi.read_states(self.states_file)
            ndstates=vaex.array_types.to_numpy(states)

        #load pf
        pf=exomolapi.read_pf(self.pf_file)
        self.gQT=jnp.array(pf["QT"].to_numpy()) #grid QT
        self.T_gQT=jnp.array(pf["T"].to_numpy()) #T forgrid QT

        self.Tref=296.0
        self.QTref=np.array(self.QT_interp(self.Tref))

        #trans file(s)
        print("Reading transition file")
        if numinf is None:
            self.trans_file = self.path/pathlib.Path(molec+".trans.bz2")
            if not self.trans_file.exists():
                self.download(molec,[".trans.bz2"])

            if self.trans_file.with_suffix(".hdf5").exists():
                trans=vaex.open(self.trans_file.with_suffix(".hdf5"))
                cdt=(trans.nu_lines>self.nurange[0]-self.margin) \
                    * (trans.nu_lines<self.nurange[1]+self.margin)
                if not np.isneginf(self.crit):
                    cdt=cdt * (trans.Sij0>self.crit)
                trans=trans[cdt]
                ndtrans=vaex.array_types.to_numpy(trans)

                #mask has been alraedy applied
                mask_needed=False
            else:
                print(explanation_trans)
                trans=exomolapi.read_trans(self.trans_file)
                ndtrans=vaex.array_types.to_numpy(trans)

                #mask needs to be applied
                mask_needed=True

            #compute gup and elower
            self._A, self.nu_lines, self._elower, self._gpp, self._jlower, self._jupper, mask_zeronu=exomolapi.pickup_gE(ndstates,ndtrans,self.trans_file)

            if self.trans_file.with_suffix(".hdf5").exists():
                self.Sij0=ndtrans[:,4]
            else:
                ##Line strength: input should be ndarray not jnp array
                self.Sij0=exomol.Sij0(self._A,self._gpp,self.nu_lines,self._elower,self.QTref)

                #exclude the lines whose nu_lines evaluated inside exomolapi.pickup_gE (thus sometimes different from the "nu_lines" column in trans) is not positive
                trans["nu_positive"]=mask_zeronu
                trans=trans[trans.nu_positive].extract()
                trans.drop('nu_positive',inplace=True)

                trans["nu_lines"]=self.nu_lines
                trans["Sij0"]=self.Sij0
                trans.export(self.trans_file.with_suffix(".hdf5"))
        else:
            imin=np.searchsorted(numinf,self.nurange[0],side="right")-1 #left side
            imax=np.searchsorted(numinf,self.nurange[1],side="right")-1 #left side
            self.trans_file=[]
            for k,i in enumerate(range(imin,imax+1)):
                trans_file = self.path/pathlib.Path(molec+"__"+numtag[i]+".trans.bz2")
                if not trans_file.exists():
                    self.download(molec,extension=[".trans.bz2"],numtag=numtag[i])

                if trans_file.with_suffix(".hdf5").exists():
                    trans=vaex.open(trans_file.with_suffix(".hdf5"))
                    cdt=(trans.nu_lines>self.nurange[0]-self.margin) \
                        * (trans.nu_lines<self.nurange[1]+self.margin)
                    if not np.isneginf(self.crit):
                        cdt=cdt * (trans.Sij0>self.crit)
                    trans=trans[cdt]
                    ndtrans=vaex.array_types.to_numpy(trans)
                    self.trans_file.append(trans_file)

                    #mask has been alraedy applied
                    mask_needed=False
                else:
                    print(explanation_trans)
                    trans=exomolapi.read_trans(trans_file)
                    ndtrans=vaex.array_types.to_numpy(trans)
                    self.trans_file.append(trans_file)

                    #mask needs to be applied
                    mask_needed=True

                #compute gup and elower
                if k==0:
                    self._A, self.nu_lines, self._elower, self._gpp, self._jlower, self._jupper, mask_zeronu=exomolapi.pickup_gE(ndstates,ndtrans,trans_file)
                    if trans_file.with_suffix(".hdf5").exists():
                        self.Sij0=ndtrans[:,4]
                    else:
                        ##Line strength: input should be ndarray not jnp array
                        self.Sij0=exomol.Sij0(self._A,self._gpp,self.nu_lines,self._elower,self.QTref)

                        #exclude the lines whose nu_lines evaluated inside exomolapi.pickup_gE (thus sometimes different from the "nu_lines" column in trans) is not positive
                        trans["nu_positive"]=mask_zeronu
                        trans=trans[trans.nu_positive].extract()
                        trans.drop('nu_positive',inplace=True)

                        trans["nu_lines"]=self.nu_lines
                        trans["Sij0"]=self.Sij0
                else:
                    Ax, nulx, elowerx, gppx, jlowerx, jupperx, mask_zeronu=exomolapi.pickup_gE(ndstates,ndtrans,trans_file)
                    if trans_file.with_suffix(".hdf5").exists():
                        Sij0x=ndtrans[:,4]
                    else:
                        ##Line strength: input should be ndarray not jnp array
                        Sij0x=exomol.Sij0(Ax,gppx,nulx,elowerx,self.QTref)

                        #exclude the lines whose nu_lines evaluated inside exomolapi.pickup_gE (thus sometimes different from the "nu_lines" column in trans) is not positive
                        trans["nu_positive"]=mask_zeronu
                        trans=trans[trans.nu_positive].extract()
                        trans.drop('nu_positive',inplace=True)

                        trans["nu_lines"]=nulx
                        trans["Sij0"]=Sij0x

                    self._A=np.hstack([self._A,Ax])
                    self.nu_lines=np.hstack([self.nu_lines,nulx])
                    self._elower=np.hstack([self._elower,elowerx])
                    self._gpp=np.hstack([self._gpp,gppx])
                    self._jlower=np.hstack([self._jlower,jlowerx])
                    self._jupper=np.hstack([self._jupper,jupperx])
                    self.Sij0=np.hstack([self.Sij0,Sij0x])

                if not trans_file.with_suffix(".hdf5").exists():
                    trans.export(trans_file.with_suffix(".hdf5"))

        ### MASKING ###
        mask=(self.nu_lines>self.nurange[0]-self.margin)\
        *(self.nu_lines<self.nurange[1]+self.margin)\
        *(self.Sij0>self.crit)

        self.masking(mask,mask_needed)

    def masking(self,mask,mask_needed=True):
        """applying mask and (re)generate jnp.arrays

        Args:
           mask: mask to be applied. self.mask is updated.
           mask_needed: whether mask needs to be applied or not

        Note:
           We have nd arrays and jnp arrays. We apply the mask to nd arrays and generate jnp array from the corresponding nd array. For instance, self._A is nd array and self.A is jnp array.

        """
        if mask_needed:
            #numpy float 64 Do not convert them jnp array
            self.nu_lines = self.nu_lines[mask]
            self.Sij0 = self.Sij0[mask]
            self._A=self._A[mask]
            self._elower=self._elower[mask]
            self._gpp=self._gpp[mask]
            self._jlower=self._jlower[mask]
            self._jupper=self._jupper[mask]

        #jnp arrays
        self.dev_nu_lines=jnp.array(self.nu_lines)
        self.logsij0=jnp.array(np.log(self.Sij0))
        self.A=jnp.array(self._A)
        self.gamma_natural=gn(self.A)
        self.elower=jnp.array(self._elower)
        self.gpp=jnp.array(self._gpp)
        self.jlower=jnp.array(self._jlower,dtype=int)
        self.jupper=jnp.array(self._jupper,dtype=int)
        ##Broadening parameters
        self.set_broadening()

    def set_broadening(self,alpha_ref_def=None,n_Texp_def=None):
        """setting broadening parameters

        Args:
           alpha_ref: set default alpha_ref and apply it. None=use self.alpha_ref_def
           n_Texp_def: set default n_Texp and apply it. None=use self.n_Texp_def

        """
        if alpha_ref_def:
            self.alpha_ref_def = alpha_ref_def
        if n_Texp_def:
            self.n_Texp_def = n_Texp_def

        if self.broadf:
            try:
                print(".broad is used.")
                bdat=exomolapi.read_broad(self.broad_file)
                codelv=exomolapi.check_bdat(bdat)
                print("Broadening code level=",codelv)
                if codelv=="a0":
                    j2alpha_ref, j2n_Texp = exomolapi.make_j2b(bdat,\
                        alpha_ref_default=self.alpha_ref_def,\
                        n_Texp_default=self.n_Texp_def,\
                        jlower_max=np.max(self._jlower))
                    self.alpha_ref=jnp.array(j2alpha_ref[self._jlower])
                    self.n_Texp=jnp.array(j2n_Texp[self._jlower])
                elif codelv=="a1":
                    j2alpha_ref, j2n_Texp = exomolapi.make_j2b(bdat,\
                        alpha_ref_default=self.alpha_ref_def,\
                        n_Texp_default=self.n_Texp_def,\
                        jlower_max=np.max(self._jlower))
                    jj2alpha_ref, jj2n_Texp=exomolapi.make_jj2b(bdat,\
                        j2alpha_ref_def=j2alpha_ref,j2n_Texp_def=j2n_Texp,\
                        jupper_max=np.max(self._jupper))
                    self.alpha_ref=jnp.array(jj2alpha_ref[self._jlower,self._jupper])
                    self.n_Texp=jnp.array(jj2n_Texp[self._jlower,self._jupper])
            except:
                print("Warning: Cannot load .broad. The default broadening parameters are used.")
                self.alpha_ref=jnp.array(self.alpha_ref_def*np.ones_like(self._jlower))
                self.n_Texp=jnp.array(self.n_Texp_def*np.ones_like(self._jlower))

        else:
            print("The default broadening parameters are used.")
            self.alpha_ref=jnp.array(self.alpha_ref_def*np.ones_like(self._jlower))
            self.n_Texp=jnp.array(self.n_Texp_def*np.ones_like(self._jlower))

    def QT_interp(self,T):
        """interpolated partition function

        Args:
           T: temperature

        Returns:
           Q(T) interpolated in jnp.array

        """
        return jnp.interp(T,self.T_gQT,self.gQT)

    def qr_interp(self,T):
        """interpolated partition function ratio

        Args:
           T: temperature

        Returns:
           qr(T)=Q(T)/Q(Tref) interpolated in jnp.array

        """
        return self.QT_interp(T)/self.QT_interp(self.Tref)

    def download(self,molec,extension,numtag=None):
        """Downloading Exomol files

        Args:
           molec: like "12C-16O__Li2015"
           extension: extension list e.g. [".pf",".def",".trans.bz2",".states.bz2",".broad"]
           numtag: number tag of transition file if exists. e.g. "11100-11200"

        Note:
           The download URL is written in exojax.utils.url.

        """
        import urllib.request
        from exojax.utils.molname import e2s
        import os
        from exojax.utils.url import url_ExoMol

        tag=molec.split("__")
        molname_simple=e2s(tag[0])

        for ext in extension:
            if ext==".trans.bz2" and numtag is not None:
                ext="__"+numtag+ext

            if ext==".broad":
                pfname_arr=[tag[0]+"__H2"+ext,tag[0]+"__He"+ext,tag[0]+"__air"+ext]
                url = url_ExoMol()+molname_simple+"/"+tag[0]+"/"
            else:
                pfname_arr=[molec+ext]
                url = url_ExoMol()+molname_simple+"/"+tag[0]+"/"+tag[1]+"/"

            for pfname in pfname_arr:
                pfpath=url+pfname
                os.makedirs(str(self.path), exist_ok=True)
                print("Downloading "+pfpath)
                try:
                    urllib.request.urlretrieve(pfpath,str(self.path/pfname))
                except:
                    print("Error: Couldn't download "+ext+" file and save.")

class MdbHit(object):
    """ molecular database of ExoMol

    MdbExomol is a class for ExoMol.

    Attributes:
        nurange: nu range [min,max] (cm-1)
        nu_lines (nd array): line center (cm-1)
        Sij0 (nd array): line strength at T=Tref (cm)
        dev_nu_lines (jnp array): line center in device (cm-1)
        logsij0 (jnp array): log line strength at T=Tref
        A (jnp array): Einstein A coeeficient
        gamma_natural (jnp array): gamma factor of the natural broadening
        gamma_air (jnp array): gamma factor of air pressure broadening
        gamma_self (jnp array): gamma factor of self pressure broadening
        elower (jnp array): the lower state energy (cm-1)
        gpp (jnp array): statistical weight
        n_air (jnp array): air temperature exponent

    """

    def __init__(self,path,nurange=[-np.inf,np.inf],margin=0.0,crit=-np.inf,extract=False):
        """Molecular database for HITRAN/HITEMP form

        Args:
           path: path for HITRAN/HITEMP par file
           nurange: wavenumber range list (cm-1) [min,max] or wavenumber grid
           margin: margin for nurange (cm-1)
           crit: line strength lower limit for extraction
           extract: If True, it extracts the opacity having the wavenumber between nurange +- margin. Use when you want to reduce the memory use.

        """
        #downloading
        self.path = pathlib.Path(path)
        if not self.path.exists():
            self.download()

        #extract?
        if extract:
            if self.path.suffix==".bz2":
                tag=str(nurange[0])+"_"+str(nurange[-1])+"_"+str(margin)
                self.path = hitranapi.extract_hitemp(str(self.path),nurange,margin,tag)
                print("self.path changed:",self.path)
            else:
                print('Warning: "extract" option is available only for .bz2 format. No "extract" applied')

        #bunzip2 if suffix is .bz2
        if self.path.suffix==".bz2":
            import bz2,shutil
            if self.path.with_suffix('').exists():
                import os
                os.remove(self.path.with_suffix(''))
            print("bunziping")
            with bz2.BZ2File(str(self.path)) as fr:
                with open(str(self.path.with_suffix('')),"wb") as fw:
                    shutil.copyfileobj(fr,fw)
            self.path=self.path.with_suffix('')

        hapi.db_begin(str(self.path.parent))
        molec=str(self.path.stem)
        self.Tref=296.0
        self.molecid = search_molecid(molec)
        self.crit = crit
        self.margin = margin
        self.nurange=[np.min(nurange),np.max(nurange)]

        #nd arrays using DRAM (not jnp, not in GPU)
        self.nu_lines = hapi.getColumn(molec, 'nu')
        self.Sij0 = hapi.getColumn(molec, 'sw')
        self.delta_air = hapi.getColumn(molec, 'delta_air')
        self.isoid = hapi.getColumn(molec,'local_iso_id')
        self.uniqiso=np.unique(self.isoid)

        self._A=hapi.getColumn(molec, 'a')
        self._n_air = hapi.getColumn(molec, 'n_air')
        self._gamma_air = hapi.getColumn(molec, 'gamma_air')
        self._gamma_self =hapi.getColumn(molec, 'gamma_self')
        self._elower = hapi.getColumn(molec, 'elower')
        self._gpp = hapi.getColumn(molec, 'gpp')

        ### MASKING ###
        mask=(self.nu_lines>self.nurange[0]-self.margin)\
        *(self.nu_lines<self.nurange[1]+self.margin)\
        *(self.Sij0>self.crit)

        self.masking(mask)

    def masking(self,mask):
        """applying mask and (re)generate jnp.arrays

        Args:
           mask: mask to be applied

        Note:
           We have nd arrays and jnp arrays. We apply the mask to nd arrays and generate jnp array from the corresponding nd array. For instance, self._A is nd array and self.A is jnp array.

        """

        #numpy float 64 Do not convert them jnp array
        self.nu_lines = self.nu_lines[mask]
        self.Sij0 = self.Sij0[mask]
        self.delta_air=self.delta_air[mask]
        self.isoid = self.isoid[mask]
        self.uniqiso=np.unique(self.isoid)

        ##numpy float 64 copy source for jnp
        self._A=self._A[mask]
        self._n_air = self._n_air[mask]
        self._gamma_air = self._gamma_air[mask]
        self._gamma_self = self._gamma_self[mask]
        self._elower = self._elower[mask]
        self._gpp = self._gpp[mask]

        #jnp.array copy from the copy sources
        self.dev_nu_lines=jnp.array(self.nu_lines)
        self.logsij0=jnp.array(np.log(self.Sij0))
        self.A=jnp.array(self._A)
        self.n_air=jnp.array(self._n_air)
        self.gamma_air = jnp.array(self._gamma_air)
        self.gamma_self = jnp.array(self._gamma_self)
        self.elower=jnp.array(self._elower)
        self.gpp=jnp.array(self._gpp)
        self.gamma_natural=gn(self.A)

    def download(self):
        """Downloading HITRAN/HITEMP par file

        Note:
           The download URL is written in exojax.utils.url.

        """
        import urllib.request
        from exojax.utils.url import url_HITRAN12
        from exojax.utils.url import url_HITEMP

        try:
            url = url_HITRAN12()+self.path.name
            urllib.request.urlretrieve(url,str(self.path))
        except:
            print(url)
            print("HITRAN download failed")
        try:
            url = url_HITEMP()+self.path.name
            print(url)
            urllib.request.urlretrieve(url,str(self.path))
        except:
            print("HITEMP download failed")

    ####################################

    def ExomolQT(self,path):
        """use a partition function from ExoMol

        Args:
           path: path for Exomol data directory/tag. For instance, "/home/CO/12C-16O/Li2015"

        """
        #load pf

        self.empath = pathlib.Path(path)
        t0=self.empath.parents[0].stem
        molec=t0+"__"+str(self.empath.stem)
        self.pf_file = self.empath/pathlib.Path(molec+".pf")
        if not self.pf_file.exists():
                self.exomol_pf_download(molec)

        pf=exomolapi.read_pf(self.pf_file)
        self.gQT=jnp.array(pf["QT"].to_numpy()) #grid QT
        self.T_gQT=jnp.array(pf["T"].to_numpy()) #T forgrid QT

    def exomol_pf_download(self,molec):
        """Downloading Exomol pf files

        Args:
           molec: like "12C-16O__Li2015"

        Note:
           The download URL is written in exojax.utils.url.

        """
        import urllib.request
        from exojax.utils.molname import e2s
        import os
        from exojax.utils.url import url_ExoMol

        tag=molec.split("__")
        molname_simple=e2s(tag[0])
        url = url_ExoMol()+molname_simple+"/"+tag[0]+"/"+tag[1]+"/"

        ext=".pf"
        pfname=molec+ext
        pfpath=url+pfname
        os.makedirs(str(self.empath), exist_ok=True)
        print("Downloading "+pfpath)
        try:
            urllib.request.urlretrieve(pfpath,str(self.empath/pfname))
        except:
            print("Error: Couldn't download "+ext+" file and save.")

    def QT_interp(self,T):
        """interpolated partition function

        Args:
           T: temperature

        Returns:
           Q(T) interpolated in jnp.array

        """
        return jnp.interp(T,self.T_gQT,self.gQT)

    def qr_interp(self,T):
        """interpolated partition function ratio

        Args:
           T: temperature

        Returns:
           qr(T)=Q(T)/Q(Tref) interpolated in jnp.array

        """
        return self.QT_interp(T)/self.QT_interp(self.Tref)

    def Qr_HAPI(self,Tarr):
        """Partition Function ratio using HAPI partition sum

        Args:
           Tarr: temperature array (K)

        Returns:
           Qr = partition function ratio array [N_Tarr x N_iso]

        Note:
           N_Tarr = len(Tarr), N_iso = len(self.uniqiso)

        """
        allT=list(np.concatenate([[self.Tref],Tarr]))
        Qrx=[]
        for iso in self.uniqiso:
            Qrx.append(hapi.partitionSum(self.molecid,iso, allT))
        Qrx=np.array(Qrx)
        qr=Qrx[:,1:].T/Qrx[:,0] #Q(T)/Q(Tref)
        return qr

    def Qr_line_HAPI(self,T):
        """Partition Function ratio using HAPI partition sum

        Args:
           T: temperature (K)

        Returns:
           Qr_line, partition function ratio array for lines [Nlines]

        Note:
           Nlines=len(self.nu_lines)

        """
        qr_line=np.ones_like(self.isoid,dtype=np.float64)
        qrx=self.Qr_HAPI([T])
        for idx,iso in enumerate(self.uniqiso):
            mask=self.isoid==iso
            qr_line[mask]=qrx[0,idx]
        return qr_line

    def Qr_layer_HAPI(self,Tarr):
        """Partition Function ratio using HAPI partition sum

        Args:
           Tarr: temperature array (K)

        Returns:
           Qr_layer, partition function ratio array for lines [N_Tarr x Nlines]

        Note:
           Nlines=len(self.nu_lines)
           N_Tarr=len(Tarr)

        """
        NP=len(Tarr)
        qt=np.zeros((NP,len(self.isoid)))
        qr=self.Qr_HAPI(Tarr)
        for idx,iso in enumerate(self.uniqiso):
            mask=self.isoid==iso
            for ilayer in range(NP):
                qt[ilayer,mask]=qr[ilayer,idx]
        return qt

def search_molecid(molec):
    """molec id from molec (source table name) of HITRAN/HITEMP

    Args:
       molec: source table name

    Return:
       int: molecid (HITRAN molecular id)

    """
    try:
        hitf=molec.split("_")
        molecid=int(hitf[0])
        return molecid

    except:
        print("Warning: Define molecid by yourself.")
        return None

if __name__ == "__main__":
    #mdb=MdbExomol("/home/kawahara/exojax/data/CO/12C-16O/Li2015/")
    #mdb=MdbExomol("/home/kawahara/exojax/data/CH4/12C-1H4/YT34to10/",nurange=[6050.0,6150.0])
    mdb=MdbExomol('.database/H2O/1H2-16O/POKAZATEL',[4310.0,4320.0],crit=1.e-45)

#    mask=mdb.A>1.e-42
#    mdb.masking(mask)
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/NH3/14N-1H3/CoYuTe/",nurange=[6050.0,6150.0])
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/H2S/1H2-32S/AYT2/",nurange=[6050.0,6150.0])
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/FeH/56Fe-1H/MoLLIST/",nurange=[6050.0,6150.0])
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/NO/14N-16O/NOname/14N-16O__NOname")

class AdbVald(object):  #integrated from vald3db.py
    """ atomic database from VALD3 (http://vald.astro.uu.se/)

    AdbVald is a class for VALD3.

    Attributes:
        nurange: nu range [min,max] (cm-1)
        nu_lines (nd array):      line center (cm-1) (#NOT frequency in (s-1))
        dev_nu_lines (jnp array): line center (cm-1) in device
        Sij0 (nd array): line strength at T=Tref (cm)
        logsij0 (jnp array): log line strength at T=Tref
        A (jnp array): Einstein A coeeficient in (s-1)
        elower (jnp array): the lower state energy (cm-1)
        eupper (jnp array): the upper state energy (cm-1)
        gupper: (jnp array): upper statistical weight
        jlower (jnp array): lower J (rotational quantum number, total angular momentum)
        jupper (jnp array): upper J
        QTmask (jnp array): identifier of species for Q(T)
        ielem (jnp array):  atomic number (e.g., Fe=26)
        iion (jnp array):  ionized level (e.g., neutral=1, singly ionized=2, etc.)
        gamRad (jnp array): log of gamma of radiation damping (s-1) #(https://www.astro.uu.se/valdwiki/Vald3Format)
        gamSta (jnp array): log of gamma of Stark damping (s-1)
        vdWdamp (jnp array):  log of (van der Waals damping constant / neutral hydrogen number) (s-1)

        Note:
           For the first time to read the VALD line list, it is converted to HDF/vaex. After the second-time, we use the HDF5 format with vaex instead.

    """
    def __init__(self, path, nurange=[-np.inf,np.inf], margin=0.0, crit=-np.inf, Irwin=False):

        """Atomic database for VALD3 "Long format"

        Args:
          path: path for linelists downloaded from VALD3 with a query of "Long format" in the format of "Extract All" and "Extract Element" (NOT "Extract Stellar")
          nurange: wavenumber range list (cm-1) or wavenumber array
          margin: margin for nurange (cm-1)
          crit: line strength lower limit for extraction
          Irwin: if True(1), the partition functions of Irwin1981 is used, otherwise those of Barklem&Collet2016

        Note:
          (written with reference to moldb.py, but without using feather format)
        """

        #load args
        self.vald3_file = pathlib.Path(path) #VALD3 output
        #self.path = pathlib.Path(path) #molec=t0+"__"+str(self.path.stem) #t0=self.path.parents[0].stem
        self.nurange = [np.min(nurange),np.max(nurange)]
        self.margin = margin
        self.crit = crit
        #self.bkgdatm=bkgdatm
        #self.broadf=broadf

        #load vald file ("Extract Stellar" request)
        print("Reading VALD file")
        if self.vald3_file.with_suffix(".hdf5").exists():
            valdd = vaex.open(self.vald3_file.with_suffix(".hdf5"))
        else:
            print("Note: Couldn't find the hdf5 format. We convert data to the hdf5 format.")
            valdd = atomllapi.read_ExAll(self.vald3_file) #vaex.DataFrame
        pvaldd = valdd.to_pandas_df() #pandas.DataFrame

        #compute additional transition parameters
        self._A, self.nu_lines, self._elower, self._eupper, self._gupper, self._jlower, self._jupper, self._ielem, self._iion, self._gamRad, self._gamSta, self._vdWdamp = atomllapi.pickup_param(pvaldd)

        #load the partition functions (for 284 atomic species)
        pfTdat, self.pfdat = atomllapi.load_pf_Barklem2016() #Barklem & Collet (2016)
        self.T_gQT = jnp.array(pfTdat.columns[1:], dtype=float)
        self.gQT_284species = jnp.array(self.pfdat.iloc[:, 1:].to_numpy(dtype=float)) #grid Q vs T vs Species
        self.Tref=296.0 #\\\\
        self.QTref_284 = np.array(self.QT_interp_284(self.Tref))
        self._QTmask = self.make_QTmask(self._ielem, self._iion) #identify index of QT grid (gQT) for each line

        ##Line strength: input shoud be ndarray not jnp array
        self.Sij0 = atomll.Sij0(self._A, self._gupper, self.nu_lines, self._elower, self.QTref_284, self._QTmask, Irwin) #211013

        ### MASKING ###
        mask=(self.nu_lines>self.nurange[0]-self.margin)\
        *(self.nu_lines<self.nurange[1]+self.margin)\
        *(self.Sij0>self.crit)

        self.masking(mask)

        #Compile atomic-specific data for each absorption line of interest
        self.ipccd = atomllapi.load_atomicdata()
        #print(self.ipccd)#test
        # should be refined
        ionE = jnp.array(list(map(lambda x: self.ipccd[self.ipccd['ielem']==x].iat[0, 1], self.ielem)))
        ionE2 = jnp.array(list(map(lambda x: self.ipccd[self.ipccd['ielem']==x].iat[0, 6], self.ielem)))
        self.ionE = ionE * np.where(self.iion==1, 1, 0) + ionE2 * np.where(self.iion==2, 1, 0)
        self.solarA = jnp.array(list(map(lambda x: self.ipccd[self.ipccd['ielem']==x].iat[0, 4], self.ielem)))
        self.atomicmass = jnp.array(list(map(lambda x: self.ipccd[self.ipccd['ielem']==x].iat[0, 5], self.ielem)))

    #End of the CONSTRUCTOR definition ↑

    #Defining METHODS ↓

    def masking(self,mask):
        """applying mask and (re)generate jnp.arrays

        Args:
           mask: mask to be applied. self.mask is updated.

        Note:
           We have nd arrays and jnp arrays. We apply the mask to nd arrays and generate jnp array from the corresponding nd array. For instance, self._A is nd array and self.A is jnp array.

        """
        #numpy float 64 Do not convert them jnp array
        self.nu_lines = self.nu_lines[mask]
        self.Sij0 = self.Sij0[mask]
        self._A=self._A[mask]
        self._elower=self._elower[mask]
        self._eupper=self._eupper[mask]
        self._gupper=self._gupper[mask]
        self._jlower=self._jlower[mask]
        self._jupper=self._jupper[mask]
        self._QTmask=self._QTmask[mask]
        self._ielem=self._ielem[mask]
        self._iion=self._iion[mask]
        self._gamRad=self._gamRad[mask]
        self._gamSta=self._gamSta[mask]
        self._vdWdamp=self._vdWdamp[mask]

        #jnp arrays
        self.dev_nu_lines=jnp.array(self.nu_lines)
        self.logsij0=jnp.array(np.log(self.Sij0))
        self.A=jnp.array(self._A)
        self.elower=jnp.array(self._elower)
        self.eupper=jnp.array(self._eupper)
        self.gupper=jnp.array(self._gupper)
        self.jlower=jnp.array(self._jlower,dtype=int)
        self.jupper=jnp.array(self._jupper,dtype=int)

        self.QTmask=jnp.array(self._QTmask,dtype=int)
        self.ielem=jnp.array(self._ielem,dtype=int)
        self.iion=jnp.array(self._iion,dtype=int)
        self.gamRad=jnp.array(self._gamRad)
        self.gamSta=jnp.array(self._gamSta)
        self.vdWdamp=jnp.array(self._vdWdamp)

    def Atomic_gQT(self, atomspecies):
        """Select grid of partition function especially for the species of interest

        Args:
            atomspecies: species e.g., "Fe 1", "Sr 2", etc.

        Returns:
            gQT: grid Q(T) for the species

        """
        #if len(atomspecies.split(' '))==2:
        atomspecies_Roman = atomspecies.split(' ')[0] + '_' + 'I'*int(atomspecies.split(' ')[-1])
        gQT = self.gQT_284species[ np.where(self.pfdat['T[K]']==atomspecies_Roman) ][0]
        return gQT

    def QT_interp(self, atomspecies, T):
        """interpolated partition function
            The partition functions of Barklem & Collet (2016) are adopted.

        Args:
          atomspecies: species e.g., "Fe 1"
          T: temperature

        Returns:
          Q(T): interpolated in jnp.array for the Atomic Species

        """
        gQT = self.Atomic_gQT(atomspecies)
        QT = jnp.interp(T, self.T_gQT, gQT)
        return QT

    def QT_interp_Irwin_Fe(self, T, atomspecies="Fe 1"):
        """interpolated partition function
            This function is for the exceptional case where you want to adopt partition functions of Irwin (1981) for Fe I (Other species are not yet implemented).

        Args:
          atomspecies: species e.g., "Fe 1"
          T: temperature

        Returns:
          Q(T): interpolated in jnp.array for the Atomic Species

        """
        gQT = self.Atomic_gQT(atomspecies)
        QT = atomllapi.partfn_Fe(T)
        return QT

    def qr_interp(self, atomspecies, T):
        """interpolated partition function ratio
            The partition functions of Barklem & Collet (2016) are adopted.

        Args:
           T: temperature
           atomspecies: species e.g., "Fe 1"

        Returns:
           qr(T)=Q(T)/Q(Tref): interpolated in jnp.array

        """
        return self.QT_interp(atomspecies,T)/self.QT_interp(atomspecies,self.Tref)

    def qr_interp_Irwin_Fe(self, T, atomspecies="Fe 1"):
        """interpolated partition function ratio
            This function is for the exceptional case where you want to adopt partition functions of Irwin (1981) for Fe I (Other species are not yet implemented).

        Args:
           T: temperature
           atomspecies: species e.g., "Fe 1"

        Returns:
           qr(T)=Q(T)/Q(Tref): interpolated in jnp.array

        """
        return self.QT_interp_Irwin_Fe(T,atomspecies)/self.QT_interp_Irwin_Fe(self.Tref,atomspecies)

    def QT_interp_284(self, T):
        """interpolated partition function of all 284 species

        Args:
           T: temperature

        Returns:
           Q(T)*284: interpolated in jnp.array for all 284 Atomic Species

        """
        #self.T_gQT.shape -> (42,)
        #self.gQT_284species.shape -> (284, 42)
        list_gQT_eachspecies = self.gQT_284species.tolist()
        listofDA_gQT_eachspecies = list(map(lambda x: jnp.array(x), list_gQT_eachspecies))
        listofQT = list(map(lambda x: jnp.interp(T, self.T_gQT, x), listofDA_gQT_eachspecies))
        QT_284 = jnp.array(listofQT)
        return QT_284

    def make_QTmask(self, ielem, iion):
        """Convert the species identifier to the index for Q(Tref) grid (gQT) for each line

        Args:
            ielem:  atomic number (e.g., Fe=26)
            iion:  ionized level (e.g., neutral=1, singly)

        Returns:
            QTmask_sp:  array of index of Q(Tref) grid (gQT) for each line

        """
        def species_to_QTmask(ielem, iion):
            sp_Roman = atomllapi.PeriodicTable[ielem] + '_' + 'I'*iion
            QTmask = np.where(self.pfdat['T[K]']==sp_Roman)[0][0]
            return QTmask
        QTmask_sp = np.array(list(map(species_to_QTmask, ielem, iion))).astype('int')
        return(QTmask_sp)

class AdbKurucz(object):
    """ atomic database from Kurucz (http://kurucz.harvard.edu/linelists/)

    AdbKurucz is a class for Kurucz line list.

    Attributes:
        nurange: nu range [min,max] (cm-1)
        nu_lines (nd array):      line center (cm-1) (#NOT frequency in (s-1))
        dev_nu_lines (jnp array): line center (cm-1) in device
        Sij0 (nd array): line strength at T=Tref (cm)
        logsij0 (jnp array): log line strength at T=Tref
        A (jnp array): Einstein A coeeficient in (s-1)
        elower (jnp array): the lower state energy (cm-1)
        eupper (jnp array): the upper state energy (cm-1)
        gupper: (jnp array): upper statistical weight
        jlower (jnp array): lower J (rotational quantum number, total angular momentum)
        jupper (jnp array): upper J
        QTmask (jnp array): identifier of species for Q(T)
        ielem (jnp array):  atomic number (e.g., Fe=26)
        iion (jnp array):  ionized level (e.g., neutral=1, singly ionized=2, etc.)
        gamRad (jnp array): log of gamma of radiation damping (s-1) #(https://www.astro.uu.se/valdwiki/Vald3Format)
        gamSta (jnp array): log of gamma of Stark damping (s-1)
        vdWdamp (jnp array):  log of (van der Waals damping constant / neutral hydrogen number) (s-1)

    """
    def __init__(self, path, nurange=[-np.inf,np.inf], margin=0.0, crit=-np.inf, Irwin=False):

        """Atomic database for Kurucz line list "gf????.all"

        Args:
          path: path for linelists (gf????.all) downloaded from the Kurucz web page
          nurange: wavenumber range list (cm-1) or wavenumber array
          margin: margin for nurange (cm-1)
          crit: line strength lower limit for extraction
          Irwin: if True(1), the partition functions of Irwin1981 is used, otherwise those of Barklem&Collet2016

        Note:
          (written with reference to moldb.py, but without using feather format)
        """

        #load args
        self.kurucz_file = pathlib.Path(path) #VALD3 output
        #self.path = pathlib.Path(path) #molec=t0+"__"+str(self.path.stem) #t0=self.path.parents[0].stem
        self.nurange = [np.min(nurange),np.max(nurange)]
        self.margin = margin
        self.crit = crit
        #self.bkgdatm=bkgdatm
        #self.broadf=broadf

        #load vald file ("Extract Stellar" request)
        print("Reading Kurucz file")
        self._A, self.nu_lines, self._elower, self._eupper, self._gupper, self._jlower, self._jupper, self._ielem, self._iion, self._gamRad, self._gamSta, self._vdWdamp = atomllapi.read_kurucz(self.kurucz_file)

        #load the partition functions (for 284 atomic species)
        pfTdat, self.pfdat = atomllapi.load_pf_Barklem2016() #Barklem & Collet (2016)
        self.T_gQT = jnp.array(pfTdat.columns[1:], dtype=float)
        self.gQT_284species = jnp.array(self.pfdat.iloc[:, 1:].to_numpy(dtype=float)) #grid Q vs T vs Species
        self.Tref=296.0
        self.QTref_284 = np.array(self.QT_interp_284(self.Tref))
        self._QTmask = self.make_QTmask(self._ielem, self._iion) #identify index of QT grid (gQT) for each line

        ##Line strength: input shoud be ndarray not jnp array
        self.Sij0 = atomll.Sij0(self._A, self._gupper, self.nu_lines, self._elower, self.QTref_284, self._QTmask, Irwin) #211013

        ### MASKING ###
        mask=(self.nu_lines>self.nurange[0]-self.margin)\
        *(self.nu_lines<self.nurange[1]+self.margin)\
        *(self.Sij0>self.crit)

        self.masking(mask)

        #Compile atomic-specific data for each absorption line of interest
        self.ipccd = atomllapi.load_atomicdata()
        #print(self.ipccd)#test
        ionE = jnp.array(list(map(lambda x: self.ipccd[self.ipccd['ielem']==x].iat[0, 1], self.ielem)))
        ionE2 = jnp.array(list(map(lambda x: self.ipccd[self.ipccd['ielem']==x].iat[0, 6], self.ielem)))
        self.ionE = ionE * np.where(self.iion==1, 1, 0) + ionE2 * np.where(self.iion==2, 1, 0)
        self.solarA = jnp.array(list(map(lambda x: self.ipccd[self.ipccd['ielem']==x].iat[0, 4], self.ielem)))
        self.atomicmass = jnp.array(list(map(lambda x: self.ipccd[self.ipccd['ielem']==x].iat[0, 5], self.ielem)))

    #End of the CONSTRUCTOR definition ↑

    #Defining METHODS ↓

    def masking(self,mask):
        """applying mask and (re)generate jnp.arrays

        Args:
           mask: mask to be applied. self.mask is updated.

        Note:
           We have nd arrays and jnp arrays. We apply the mask to nd arrays and generate jnp array from the corresponding nd array. For instance, self._A is nd array and self.A is jnp array.

        """
        #numpy float 64 Do not convert them jnp array
        self.nu_lines = self.nu_lines[mask]
        self.Sij0 = self.Sij0[mask]
        self._A=self._A[mask]
        self._elower=self._elower[mask]
        self._eupper=self._eupper[mask]
        self._gupper=self._gupper[mask]
        self._jlower=self._jlower[mask]
        self._jupper=self._jupper[mask]
        self._QTmask=self._QTmask[mask]
        self._ielem=self._ielem[mask]
        self._iion=self._iion[mask]
        self._gamRad=self._gamRad[mask]
        self._gamSta=self._gamSta[mask]
        self._vdWdamp=self._vdWdamp[mask]

        #jnp arrays
        self.dev_nu_lines=jnp.array(self.nu_lines)
        self.logsij0=jnp.array(np.log(self.Sij0))
        self.A=jnp.array(self._A)
        self.elower=jnp.array(self._elower)
        self.eupper=jnp.array(self._eupper)
        self.gupper=jnp.array(self._gupper)
        self.jlower=jnp.array(self._jlower,dtype=int)
        self.jupper=jnp.array(self._jupper,dtype=int)

        self.QTmask=jnp.array(self._QTmask,dtype=int)
        self.ielem=jnp.array(self._ielem,dtype=int)
        self.iion=jnp.array(self._iion,dtype=int)
        self.gamRad=jnp.array(self._gamRad)
        self.gamSta=jnp.array(self._gamSta)
        self.vdWdamp=jnp.array(self._vdWdamp)

    def Atomic_gQT(self, atomspecies):
        """Select grid of partition function especially for the species of interest

        Args:
            atomspecies: species e.g., "Fe 1", "Sr 2", etc.

        Returns:
            gQT: grid Q(T) for the species

        """
        #if len(atomspecies.split(' '))==2:
        atomspecies_Roman = atomspecies.split(' ')[0] + '_' + 'I'*int(atomspecies.split(' ')[-1])
        gQT = self.gQT_284species[ np.where(self.pfdat['T[K]']==atomspecies_Roman) ][0]
        return gQT

    def QT_interp(self, atomspecies, T):
        """interpolated partition function
            The partition functions of Barklem & Collet (2016) are adopted.

        Args:
          atomspecies: species e.g., "Fe 1"
          T: temperature

        Returns:
          Q(T): interpolated in jnp.array for the Atomic Species

        """
        gQT = self.Atomic_gQT(atomspecies)
        QT = jnp.interp(T, self.T_gQT, gQT)
        return QT

    def QT_interp_Irwin_Fe(self, T, atomspecies="Fe 1"):
        """interpolated partition function
            This function is for the exceptional case where you want to adopt partition functions of Irwin (1981) for Fe I (Other species are not yet implemented).

        Args:
          atomspecies: species e.g., "Fe 1"
          T: temperature

        Returns:
          Q(T): interpolated in jnp.array for the Atomic Species

        """
        gQT = self.Atomic_gQT(atomspecies)
        QT = atomllapi.partfn_Fe(T)
        return QT

    def qr_interp(self, atomspecies, T):
        """interpolated partition function ratio
            The partition functions of Barklem & Collet (2016) are adopted.

        Args:
           T: temperature
           atomspecies: species e.g., "Fe 1"

        Returns:
           qr(T)=Q(T)/Q(Tref): interpolated in jnp.array

        """
        return self.QT_interp(atomspecies,T)/self.QT_interp(atomspecies,self.Tref)

    def qr_interp_Irwin_Fe(self, T, atomspecies="Fe 1"):
        """interpolated partition function ratio
            This function is for the exceptional case where you want to adopt partition functions of Irwin (1981) for Fe I (Other species are not yet implemented).

        Args:
           T: temperature
           atomspecies: species e.g., "Fe 1"

        Returns:
           qr(T)=Q(T)/Q(Tref): interpolated in jnp.array

        """
        return self.QT_interp_Irwin_Fe(T,atomspecies)/self.QT_interp_Irwin_Fe(self.Tref,atomspecies)

    def QT_interp_284(self, T):
        """interpolated partition function of all 284 species

        Args:
           T: temperature

        Returns:
           Q(T)*284: interpolated in jnp.array for all 284 Atomic Species

        """
        #self.T_gQT.shape -> (42,)
        #self.gQT_284species.shape -> (284, 42)
        list_gQT_eachspecies = self.gQT_284species.tolist()
        listofDA_gQT_eachspecies = list(map(lambda x: jnp.array(x), list_gQT_eachspecies))
        listofQT = list(map(lambda x: jnp.interp(T, self.T_gQT, x), listofDA_gQT_eachspecies))
        QT_284 = jnp.array(listofQT)
        return QT_284

    def make_QTmask(self, ielem, iion):
        """Convert the species identifier to the index for Q(Tref) grid (gQT) for each line

        Args:
            ielem:  atomic number (e.g., Fe=26)
            iion:  ionized level (e.g., neutral=1, singly)

        Returns:
            QTmask_sp:  array of index of Q(Tref) grid (gQT) for each line

        """
        def species_to_QTmask(ielem, iion):
            sp_Roman = atomllapi.PeriodicTable[ielem] + '_' + 'I'*iion
            QTmask = np.where(self.pfdat['T[K]']==sp_Roman)[0][0]
            return QTmask
        QTmask_sp = np.array(list(map(species_to_QTmask, ielem, iion))).astype('int')
        return(QTmask_sp)


import argparse
import glob
import os
import time
import sys

sys.path.insert(1, './nerf')
os.environ['GPU_DEBUG']='3'
import numpy as np
import torch
import torchvision
import yaml
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm, trange
import matplotlib.pyplot as plt

from nerf.load_flame import load_flame_data

from nerf import (CfgNode, get_embedding_function, get_ray_bundle, img2mse,
                  load_llff_data, meshgrid_xy, models,
                  mse2psnr, run_one_iter_of_nerf, dump_rays, GaussianSmoothing)
#from gpu_profile import gpu_profile

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config", type=str, required=True, help="Path to (.yml) config file."
    )
    parser.add_argument(
        "--load-checkpoint",
        type=str,
        default="",
        help="Path to load saved checkpoint from.",
    )
    configargs = parser.parse_args()

    # Read config file.
    cfg = None
    with open(configargs.config, "r") as f:
        cfg_dict = yaml.load(f, Loader=yaml.FullLoader)
        cfg = CfgNode(cfg_dict)

    # # (Optional:) enable this to track autograd issues when debugging
    # torch.autograd.set_detect_anomaly(True)

    # If a pre-cached dataset is available, skip the dataloader.
    USE_CACHED_DATASET = False
    train_paths, validation_paths = None, None
    images, poses, render_poses, hwf, i_split, expressions = None, None, None, None, None, None
    H, W, focal, i_train, i_val, i_test = None, None, None, None, None, None
    if hasattr(cfg.dataset, "cachedir") and os.path.exists(cfg.dataset.cachedir):
        train_paths = glob.glob(os.path.join(cfg.dataset.cachedir, "train", "*.data"))
        validation_paths = glob.glob(
            os.path.join(cfg.dataset.cachedir, "val", "*.data")
        )
        USE_CACHED_DATASET = True
    else:
        # Load dataset
        images, poses, render_poses, hwf, expressions = None, None, None, None, None
        if cfg.dataset.type.lower() == "blender":
            images, poses, render_poses, hwf, i_split, expressions, _, bboxs = load_flame_data(
                cfg.dataset.basedir,
                half_res=cfg.dataset.half_res,
                testskip=cfg.dataset.testskip,
            )
            i_train, i_val, i_test = i_split
            H, W, focal = hwf
            H, W = int(H), int(W)
            hwf = [H, W, focal]
            if cfg.nerf.train.white_background:
                images = images[..., :3] * images[..., -1:] + (1.0 - images[..., -1:])
    print("done loading data")
    # Seed experiment for repeatability
    seed = cfg.experiment.randomseed
    np.random.seed(seed)
    torch.manual_seed(seed)

    # Device on which to run.
    if torch.cuda.is_available():
        device = "cuda" #+ ":" + str(cfg.experiment.device)
    else:
        device = "cpu"

    encode_position_fn = get_embedding_function(
        num_encoding_functions=cfg.models.coarse.num_encoding_fn_xyz,
        include_input=cfg.models.coarse.include_input_xyz,
        log_sampling=cfg.models.coarse.log_sampling_xyz,
    )

    encode_direction_fn = None
    if cfg.models.coarse.use_viewdirs:
        encode_direction_fn = get_embedding_function(
            num_encoding_functions=cfg.models.coarse.num_encoding_fn_dir,
            include_input=cfg.models.coarse.include_input_dir,
            log_sampling=cfg.models.coarse.log_sampling_dir,
        )

    # Initialize a coarse-resolution model.
    model_coarse = getattr(models, cfg.models.coarse.type)(
        num_encoding_fn_xyz=cfg.models.coarse.num_encoding_fn_xyz,
        num_encoding_fn_dir=cfg.models.coarse.num_encoding_fn_dir,
        include_input_xyz=cfg.models.coarse.include_input_xyz,
        include_input_dir=cfg.models.coarse.include_input_dir,
        use_viewdirs=cfg.models.coarse.use_viewdirs,
        num_layers=cfg.models.coarse.num_layers,
        hidden_size=cfg.models.coarse.hidden_size,
        include_expression=True
    )
    model_coarse.to(device)
    # If a fine-resolution model is specified, initialize it.
    model_fine = None
    if hasattr(cfg.models, "fine"):
        model_fine = getattr(models, cfg.models.fine.type)(
            num_encoding_fn_xyz=cfg.models.fine.num_encoding_fn_xyz,
            num_encoding_fn_dir=cfg.models.fine.num_encoding_fn_dir,
            include_input_xyz=cfg.models.fine.include_input_xyz,
            include_input_dir=cfg.models.fine.include_input_dir,
            use_viewdirs=cfg.models.fine.use_viewdirs,
            num_layers = cfg.models.coarse.num_layers,
            hidden_size =cfg.models.coarse.hidden_size,
            include_expression=True
        )
        model_fine.to(device)

    ###################################
    ###################################
    train_background = False
    supervised_train_background = False
    blur_background = False

    train_latent_codes = True
    disable_expressions = False # True to disable expressions
    disable_latent_codes = False # True to disable latent codes
    fixed_background = True # Do False to disable BG
    regularize_latent_codes = True # True to add latent code LOSS, false for most experiments
    ###################################
    ###################################

    supervised_train_background = train_background and supervised_train_background
    # Avg background
    #images[i_train]
    if train_background:
        with torch.no_grad():
            avg_img = torch.mean(images[i_train],axis=0)
            # Blur Background:
            if blur_background:
                avg_img = avg_img.permute(2,0,1)
                avg_img = avg_img.unsqueeze(0)
                smoother = GaussianSmoothing(channels=3, kernel_size=11, sigma=11)
                print("smoothed background initialization. shape ", avg_img.shape)
                avg_img = smoother(avg_img).squeeze(0).permute(1,2,0)
            #avg_img = torch.zeros(H,W,3)
            #avg_img = torch.rand(H,W,3)
            #avg_img = 0.5*(torch.rand(H,W,3) + torch.mean(images[i_train],axis=0))
            background = torch.tensor(avg_img, device=device)
        background.requires_grad = True

    if fixed_background: # load GT background
        print("loading GT background to condition on")
        from PIL import Image
        background = Image.open(os.path.join(cfg.dataset.basedir,'bg','00050.png'))
        background.thumbnail((H,W))
        background = torch.from_numpy(np.array(background).astype(np.float32)).to(device)
        background = background/255
        print("bg shape", background.shape)
        print("should be ", images[i_train][0].shape)
        assert background.shape == images[i_train][0].shape
    else:
        background = None

    # Initialize optimizer.
    trainable_parameters = list(model_coarse.parameters())
    if model_fine is not None:
        trainable_parameters += list(model_fine.parameters())
    if train_background:
        #background.requires_grad = True
        #trainable_parameters.append(background) # add it later when init optimizer for different lr
        print("background.is_leaf " ,background.is_leaf, background.device)

    if train_latent_codes:
        latent_codes = torch.zeros(len(i_train),32, device=device)
        print("initialized latent codes with shape %d X %d" % (latent_codes.shape[0], latent_codes.shape[1]))
        if not disable_latent_codes:
            trainable_parameters.append(latent_codes)
            latent_codes.requires_grad = True

    if train_background:
        optimizer = getattr(torch.optim, cfg.optimizer.type)(
            [{'params':trainable_parameters},
             {'params':background, 'lr':cfg.optimizer.lr}],
            lr=cfg.optimizer.lr
        )
    else:
        optimizer = getattr(torch.optim, cfg.optimizer.type)(
            [{'params':trainable_parameters},
             {'params': background, 'lr': cfg.optimizer.lr}        ], # this is obsolete but need for continuing training
            lr=cfg.optimizer.lr
        )
    # Setup logging.
    logdir = os.path.join(cfg.experiment.logdir, cfg.experiment.id)
    os.makedirs(logdir, exist_ok=True)
    writer = SummaryWriter(logdir)
    # Write out config parameters.
    with open(os.path.join(logdir, "config.yml"), "w") as f:
        f.write(cfg.dump())  # cfg, f, default_flow_style=False)

    # By default, start at iteration 0 (unless a checkpoint is specified).
    start_iter = 0

    # Load an existing checkpoint, if a path is specified.
    if os.path.exists(configargs.load_checkpoint):
        checkpoint = torch.load(configargs.load_checkpoint)
        model_coarse.load_state_dict(checkpoint["model_coarse_state_dict"])
        if checkpoint["model_fine_state_dict"]:
            model_fine.load_state_dict(checkpoint["model_fine_state_dict"])
        if checkpoint["background"] is not None:
            print("loaded bg from checkpoint")
            background = torch.nn.Parameter(checkpoint['background'].to(device))
        if checkpoint["latent_codes"] is not None:
            print("loaded latent codes from checkpoint")
            latent_codes = torch.nn.Parameter(checkpoint['latent_codes'].to(device))

        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        start_iter = checkpoint["iter"]

    # # TODO: Prepare raybatch tensor if batching random rays

    # Prepare importance sampling maps
    ray_importance_sampling_maps = []
    p = 0.9
    print("computing boundix boxes probability maps")
    for i in i_train:
        bbox = bboxs[i]
        probs = np.zeros((H,W))
        probs.fill(1-p)
        probs[bbox[0]:bbox[1],bbox[2]:bbox[3]] = p
        probs = (1/probs.sum()) * probs
        ray_importance_sampling_maps.append(probs.reshape(-1))

    print("Starting loop")
    for i in trange(start_iter, cfg.experiment.train_iters):

        model_coarse.train()
        if model_fine:
            model_coarse.train()

        rgb_coarse, rgb_fine = None, None
        target_ray_values = None
        background_ray_values = None
        if USE_CACHED_DATASET:
            datafile = np.random.choice(train_paths)
            cache_dict = torch.load(datafile)
            ray_bundle = cache_dict["ray_bundle"].to(device)
            ray_origins, ray_directions = (
                ray_bundle[0].reshape((-1, 3)),
                ray_bundle[1].reshape((-1, 3)),
            )
            target_ray_values = cache_dict["target"][..., :3].reshape((-1, 3))
            select_inds = np.random.choice(
                ray_origins.shape[0],
                size=(cfg.nerf.train.num_random_rays),
                replace=False,
            )
            ray_origins, ray_directions = (
                ray_origins[select_inds],
                ray_directions[select_inds],
            )
            target_ray_values = target_ray_values[select_inds].to(device)
            #target_ray_values = target_ray_values[select_inds].to(device)
            # ray_bundle = torch.stack([ray_origins, ray_directions], dim=0).to(device)

            rgb_coarse, _, _, rgb_fine, _, _ = run_one_iter_of_nerf(
                cache_dict["height"],
                cache_dict["width"],
                cache_dict["focal_length"],
                model_coarse,
                model_fine,
                ray_origins,
                ray_directions,
                cfg,
                mode="train",
                encode_position_fn=encode_position_fn,
                encode_direction_fn=encode_direction_fn,
                expressions=expressions
            )
        else:
            img_idx = np.random.choice(i_train)
            img_target = images[img_idx].to(device)
            pose_target = poses[img_idx, :3, :4].to(device)
            if not disable_expressions:
                expression_target = expressions[img_idx].to(device) # vector
            else: # zero expr
                expression_target = torch.zeros(76, device=device)
            #bbox = bboxs[img_idx]
            if not disable_latent_codes:
                latent_code = latent_codes[img_idx].to(device) if train_latent_codes else None
            else:
                latent_codes = torch.zeros(32, device=device)
            #latent_code = torch.zeros(32).to(device)
            ray_origins, ray_directions = get_ray_bundle(H, W, focal, pose_target)
            coords = torch.stack(
                meshgrid_xy(torch.arange(H).to(device), torch.arange(W).to(device)),
                dim=-1,
            )

            # Only randomly choose rays that are in the bounding box !
            # coords = torch.stack(
            #     meshgrid_xy(torch.arange(bbox[0],bbox[1]).to(device), torch.arange(bbox[2],bbox[3]).to(device)),
            #     dim=-1,
            # )

            coords = coords.reshape((-1, 2))
            # select_inds = np.random.choice(
            #     coords.shape[0], size=(cfg.nerf.train.num_random_rays), replace=False
            # )

            # Use importance sampling to sample mainly in the bbox with prob p
            select_inds = np.random.choice(
                coords.shape[0], size=(cfg.nerf.train.num_random_rays), replace=False, p=ray_importance_sampling_maps[img_idx]
            )

            select_inds = coords[select_inds]
            ray_origins = ray_origins[select_inds[:, 0], select_inds[:, 1], :]
            ray_directions = ray_directions[select_inds[:, 0], select_inds[:, 1], :]
            #dump_rays(ray_origins, ray_directions)

            # batch_rays = torch.stack([ray_origins, ray_directions], dim=0)
            target_s = img_target[select_inds[:, 0], select_inds[:, 1], :]
            background_ray_values = background[select_inds[:, 0], select_inds[:, 1], :] if (train_background or fixed_background) else None
            #if i<10000:
            #   background_ray_values = None
            #background_ray_values = None
            then = time.time()
            rgb_coarse, _, _, rgb_fine, _, _, weights = run_one_iter_of_nerf(
                H,
                W,
                focal,
                model_coarse,
                model_fine,
                ray_origins,
                ray_directions,
                cfg,
                mode="train",
                encode_position_fn=encode_position_fn,
                encode_direction_fn=encode_direction_fn,
                expressions = expression_target,
                background_prior=background_ray_values,
                latent_code = latent_code if not disable_latent_codes else torch.zeros(32,device=device)

            )
            target_ray_values = target_s

        coarse_loss = torch.nn.functional.mse_loss(
            rgb_coarse[..., :3], target_ray_values[..., :3]
        )
        fine_loss = None
        if rgb_fine is not None:
            fine_loss = torch.nn.functional.mse_loss(
                rgb_fine[..., :3], target_ray_values[..., :3]
            )
        # loss = torch.nn.functional.mse_loss(rgb_pred[..., :3], target_s[..., :3])
        loss = 0.0
        # if fine_loss is not None:
        #     loss = fine_loss
        # else:
        #     loss = coarse_loss

        latent_code_loss = torch.zeros(1, device=device)
        if train_latent_codes and not disable_latent_codes:
            latent_code_loss = torch.norm(latent_code) * 0.0005
            #latent_code_loss = torch.zeros(1)

        background_loss = torch.zeros(1, device=device)
        if supervised_train_background:
            background_loss = torch.nn.functional.mse_loss(
                background_ray_values[..., :3], target_ray_values[..., :3], reduction='none'
            ).sum(1)
            background_loss = torch.mean(background_loss*weights) * 0.001

        loss = coarse_loss + (fine_loss if fine_loss is not None else 0.0)
        psnr = mse2psnr(loss.item())

        #loss_total = loss #+ (latent_code_loss if latent_code_loss is not None else 0.0)
        loss = loss + (latent_code_loss*10 if regularize_latent_codes else 0.0)
        loss_total = loss + (background_loss if supervised_train_background is not None else 0.0)
        #loss.backward()
        loss_total.backward()
        #psnr = mse2psnr(loss.item())
        optimizer.step()
        optimizer.zero_grad()

        # Learning rate updates
        num_decay_steps = cfg.scheduler.lr_decay * 1000
        lr_new = cfg.optimizer.lr * (
            cfg.scheduler.lr_decay_factor ** (i / num_decay_steps)
        )
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr_new

        if i % cfg.experiment.print_every == 0 or i == cfg.experiment.train_iters - 1:
            tqdm.write(
                "[TRAIN] Iter: "
                + str(i)
                + " Loss: "
                + str(loss.item())
                + " BG Loss: "
                + str(background_loss.item())
                + " PSNR: "
                + str(psnr)
                + " LatentReg: "
                + str(latent_code_loss.item())
            )
        #writer.add_scalar("train/loss", loss.item(), i)
        if train_latent_codes:
            writer.add_scalar("train/code_loss", latent_code_loss.item(), i)
        if supervised_train_background:
            writer.add_scalar("train/bg_loss", background_loss.item(), i)

        writer.add_scalar("train/coarse_loss", coarse_loss.item(), i)
        if rgb_fine is not None:
            writer.add_scalar("train/fine_loss", fine_loss.item(), i)
        writer.add_scalar("train/psnr", psnr, i)

        # Validation
        if (
            i % cfg.experiment.validate_every == 0
            or i == cfg.experiment.train_iters - 1 and False
        ):
            #torch.cuda.empty_cache()
            tqdm.write("[VAL] =======> Iter: " + str(i))
            model_coarse.eval()
            if model_fine:
                model_coarse.eval()

            start = time.time()
            with torch.no_grad():
                rgb_coarse, rgb_fine = None, None
                target_ray_values = None
                if USE_CACHED_DATASET:
                    datafile = np.random.choice(validation_paths)
                    cache_dict = torch.load(datafile)
                    rgb_coarse, _, _, rgb_fine, _, weights = run_one_iter_of_nerf(
                        cache_dict["height"],
                        cache_dict["width"],
                        cache_dict["focal_length"],
                        model_coarse,
                        model_fine,
                        cache_dict["ray_origins"].to(device),
                        cache_dict["ray_directions"].to(device),
                        cfg,
                        mode="validation",
                        encode_position_fn=encode_position_fn,
                        encode_direction_fn=encode_direction_fn,
                        expressions = expression_target,
                        latent_code = torch.zeros(32, device=device)
                    )
                    target_ray_values = cache_dict["target"].to(device)
                else:
                    # Do all validation set...
                    loss = 0
                    for img_idx in i_val[:2]:
                        img_target = images[img_idx].to(device)
                        #tqdm.set_description('val im %d' % img_idx)
                        #tqdm.refresh()  # to show immediately the update

                            # # save val image for debug ### DEBUG ####
                        # #GT = target_ray_values[..., :3]
                        # import PIL.Image
                        # #img = GT.permute(2, 0, 1)
                        # # Conver to PIL Image and then np.array (output shape: (H, W, 3))
                        # #im_numpy = img_target.detach().cpu().numpy()
                        # #im_numpy = np.array(torchvision.transforms.ToPILImage()(img_target.detach().cpu()))
                        #
                        # #                   im = PIL.Image.fromarray(im_numpy)
                        # im = img_target
                        # im = im.permute(2, 0, 1)
                        # img = np.array(torchvision.transforms.ToPILImage()(im.detach().cpu()))
                        # im = PIL.Image.fromarray(img)
                        # im.save('val_im_target_debug.png')
                        # ### DEBUG #### END

                        pose_target = poses[img_idx, :3, :4].to(device)
                        ray_origins, ray_directions = get_ray_bundle(
                            H, W, focal, pose_target
                        )
                        rgb_coarse, _, _, rgb_fine, _, _ ,weights= run_one_iter_of_nerf(
                            H,
                            W,
                            focal,
                            model_coarse,
                            model_fine,
                            ray_origins,
                            ray_directions,
                            cfg,
                            mode="validation",
                            encode_position_fn=encode_position_fn,
                            encode_direction_fn=encode_direction_fn,
                            expressions = expression_target,
                            background_prior = background.view(-1,3) if (train_background or fixed_background) else None,
                            latent_code = torch.zeros(32).to(device) if train_latent_codes or disable_latent_codes else None,

                        )
                        #print("did one val")
                        target_ray_values = img_target
                        coarse_loss = img2mse(rgb_coarse[..., :3], target_ray_values[..., :3])
                        curr_loss, curr_fine_loss = 0.0, 0.0
                        if rgb_fine is not None:
                            curr_fine_loss = img2mse(rgb_fine[..., :3], target_ray_values[..., :3])
                            curr_loss = curr_fine_loss
                        else:
                            curr_loss = coarse_loss
                        loss += curr_loss + curr_fine_loss

                loss /= len(i_val)
                psnr = mse2psnr(loss.item())
                writer.add_scalar("validation/loss", loss.item(), i)
                writer.add_scalar("validation/coarse_loss", coarse_loss.item(), i)
                writer.add_scalar("validation/psnr", psnr, i)
                writer.add_image(
                    "validation/rgb_coarse", cast_to_image(rgb_coarse[..., :3]), i
                )
                if rgb_fine is not None:
                    writer.add_image(
                        "validation/rgb_fine", cast_to_image(rgb_fine[..., :3]), i
                    )
                    writer.add_scalar("validation/fine_loss", fine_loss.item(), i)

                writer.add_image(
                    "validation/img_target",
                    cast_to_image(target_ray_values[..., :3]),
                    i,
                )
                if train_background or fixed_background:
                    writer.add_image(
                        "validation/background", cast_to_image(background[..., :3]), i
                    )
                    writer.add_image(
                        "validation/weights", (weights.detach().cpu().numpy()), i, dataformats='HW'
                    )
                tqdm.write(
                    "Validation loss: "
                    + str(loss.item())
                    + " Validation PSNR: "
                    + str(psnr)
                    + " Time: "
                    + str(time.time() - start)
                )

        #gpu_profile(frame=sys._getframe(), event='line', arg=None)

        if i % cfg.experiment.save_every == 0 or i == cfg.experiment.train_iters - 1:
            checkpoint_dict = {
                "iter": i,
                "model_coarse_state_dict": model_coarse.state_dict(),
                "model_fine_state_dict": None
                if not model_fine
                else model_fine.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "loss": loss,
                "psnr": psnr,
                "background": None
                if not (train_background or fixed_background)
                else background.data,
                "latent_codes": None if not train_latent_codes else latent_codes.data
            }
            torch.save(
                checkpoint_dict,
                os.path.join(logdir, "checkpoint" + str(i).zfill(5) + ".ckpt"),
            )
            tqdm.write("================== Saved Checkpoint =================")

    print("Done!")

def cast_to_image(tensor):
    # Input tensor is (H, W, 3). Convert to (3, H, W).
    tensor = tensor.permute(2, 0, 1)
    tensor = tensor.clamp(0.0,1.0)
    # Conver to PIL Image and then np.array (output shape: (H, W, 3))
    img = np.array(torchvision.transforms.ToPILImage()(tensor.detach().cpu()))
    # Map back to shape (3, H, W), as tensorboard needs channels first.
    img = np.moveaxis(img, [-1], [0])
    return img

def handle_pdb(sig, frame):
    import pdb
    pdb.Pdb().set_trace(frame)

if __name__ == "__main__":
    import signal

    print("before signal registration")
    signal.signal(signal.SIGUSR1, handle_pdb)
    print("after registration")
    #sys.settrace(gpu_profile)

    main()

"""
# Validation
        if (
            i % cfg.experiment.validate_every == 0
            or i == cfg.experiment.train_iters - 1
        ):
            tqdm.write("[VAL] =======> Iter: " + str(i))
            model_coarse.eval()
            if model_fine:
                model_coarse.eval()

            start = time.time()
            with torch.no_grad():
                rgb_coarse, rgb_fine = None, None
                target_ray_values = None
                if USE_CACHED_DATASET:
                    datafile = np.random.choice(validation_paths)
                    cache_dict = torch.load(datafile)
                    rgb_coarse, _, _, rgb_fine, _, _ = run_one_iter_of_nerf(
                        cache_dict["height"],
                        cache_dict["width"],
                        cache_dict["focal_length"],
                        model_coarse,
                        model_fine,
                        cache_dict["ray_origins"].to(device),
                        cache_dict["ray_directions"].to(device),
                        cfg,
                        mode="validation",
                        encode_position_fn=encode_position_fn,
                        encode_direction_fn=encode_direction_fn,
                    )
                    target_ray_values = cache_dict["target"].to(device)
                else:
                    img_idx = np.random.choice(i_val)
                    img_target = images[img_idx].to(device)

                    # # save val image for debug ### DEBUG ####
                    # #GT = target_ray_values[..., :3]
                    # import PIL.Image
                    # #img = GT.permute(2, 0, 1)
                    # # Conver to PIL Image and then np.array (output shape: (H, W, 3))
                    # #im_numpy = img_target.detach().cpu().numpy()
                    # #im_numpy = np.array(torchvision.transforms.ToPILImage()(img_target.detach().cpu()))
                    #
                    # #                   im = PIL.Image.fromarray(im_numpy)
                    # im = img_target
                    # im = im.permute(2, 0, 1)
                    # img = np.array(torchvision.transforms.ToPILImage()(im.detach().cpu()))
                    # im = PIL.Image.fromarray(img)
                    # im.save('val_im_target_debug.png')
                    # ### DEBUG #### END

                    pose_target = poses[img_idx, :3, :4].to(device)
                    ray_origins, ray_directions = get_ray_bundle(
                        H, W, focal, pose_target
                    )
                    rgb_coarse, _, _, rgb_fine, _, _ = run_one_iter_of_nerf(
                        H,
                        W,
                        focal,
                        model_coarse,
                        model_fine,
                        ray_origins,
                        ray_directions,
                        cfg,
                        mode="validation",
                        encode_position_fn=encode_position_fn,
                        encode_direction_fn=encode_direction_fn,
                    )
                    target_ray_values = img_target
                coarse_loss = img2mse(rgb_coarse[..., :3], target_ray_values[..., :3])
                loss, fine_loss = 0.0, 0.0
                if rgb_fine is not None:
                    fine_loss = img2mse(rgb_fine[..., :3], target_ray_values[..., :3])
                    loss = fine_loss
                else:
                    loss = coarse_loss
                loss = coarse_loss + fine_loss
                psnr = mse2psnr(loss.item())
                writer.add_scalar("validation/loss", loss.item(), i)
                writer.add_scalar("validation/coarse_loss", coarse_loss.item(), i)
                writer.add_scalar("validataion/psnr", psnr, i)
                writer.add_image(
                    "validation/rgb_coarse", cast_to_image(rgb_coarse[..., :3]), i
                )
                if rgb_fine is not None:
                    writer.add_image(
                        "validation/rgb_fine", cast_to_image(rgb_fine[..., :3]), i
                    )
                    writer.add_scalar("validation/fine_loss", fine_loss.item(), i)

                writer.add_image(
                    "validation/img_target",
                    cast_to_image(target_ray_values[..., :3]),
                    i,
                )
                tqdm.write(
                    "Validation loss: "
                    + str(loss.item())
                    + " Validation PSNR: "
                    + str(psnr)
                    + " Time: "
                    + str(time.time() - start)
                )
"""


#!/usr/bin/env python
# -*- coding: utf-8 -*-

################################################################################
#
#   MEASURE - Master Equation Automatic Solver for Unimolecular REactions
#
#   Copyright (c) 2010 by Joshua W. Allen (jwallen@mit.edu)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the 'Software'),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

"""
Contains functions for directly solving the full or reduced master equation
matrix to generate concentration profiles and population distributions. This
information is particularly useful in comparing and evaluating the methods for
reducing the master equation.
"""

import numpy
import scipy.integrate
#from pydas import DASSL

import rmgpy.constants as constants

################################################################################

def residual(t, y, K):
    return numpy.dot(K, y)

def jacobian(t, y, K):
    return K

################################################################################

def solveFullME(T, P, Elist, tlist, x0, M, indices, densStates, Nisom, Nreac, Nprod):
    """
    Directly solve the full master equation using a stiff ODE solver. Pass the
    reaction `network` to solve, the temperature `T` in K and pressure `P` in
    Pa to solve at, the energies `Elist` in J/mol to use, the output time
    points `tlist` in s, the initial total populations `x0`, the full master
    equation matrix `M`, the accounting matrix `indices` relating isomer and
    energy grain indices to indices of the master equation matrix, and the
    densities of states `densStates` in mol/J of each isomer.
    Returns the times in s, population distributions for each isomer, and total
    population profiles for each configuration.
    """

    Ngrains = len(Elist)
    Ntime = len(tlist)

    # Get equilibrium distributions
    eqDist = numpy.zeros_like(densStates)
    for i in range(Nisom):
        eqDist[i,:] = densStates[i,:] * numpy.exp(-Elist / constants.R / T)
        eqDist[i,:] /= sum(eqDist[i,:])

    # Set initial conditions
    p0 = numpy.zeros([M.shape[0]], float)
    for i in range(Nisom):
        for r in range(Ngrains):
            if indices[r,i] > 0:
                p0[indices[r,i]] = x0[i] * eqDist[i,r]
    for i in range(Nreac+Nprod):
        p0[-Nreac-Nprod + i] = x0[i+Nisom]

#    # Set up ODEs
#    me = MasterEquation(M)
#    me.initialize(t0=0, y0=p0, atol=1e-16, rtol=1e-8)
#
#    # Generate solution
#    t = numpy.zeros([Ntime], float)
#    p = numpy.zeros([Ntime, Nisom, Ngrains], float)
#    x = numpy.zeros([Ntime, Nisom+Nreac+Nprod], float)
#    for s in range(Ntime):
#        me.advance(tlist[s])
#        print me.t
#        t[s] = me.t
#        for r in range(Ngrains):
#            for i in range(0, Nisom):
#                if indices[r,i] > 0:
#                    p[s,i,r] += me.y[indices[r,i]]
#                    x[s,i] += me.y[indices[r,i]]
#        for n in range(Nisom, Nisom+Nreac+Nprod):
#            x[s,n] = me.y[-(Nisom+Nreac+Nprod)+n]

    # Set up ODEs
    ode = scipy.integrate.ode(residual, jacobian).set_integrator('vode', method='bdf', with_jacobian=True, atol=1e-16, rtol=1e-8)
    ode.set_initial_value(p0, 0.0).set_f_params(M).set_jac_params(M)

    # Generate solution
    t = numpy.zeros([Ntime], float)
    p = numpy.zeros([Ntime, Nisom, Ngrains], float)
    x = numpy.zeros([Ntime, Nisom+Nreac+Nprod], float)
    for s in range(Ntime):
        ode.integrate(tlist[s])
        t[s] = ode.t
        for r in range(Ngrains):
            for i in range(0, Nisom):
                if indices[r,i] > 0:
                    p[s,i,r] += ode.y[indices[r,i]]
                    x[s,i] += ode.y[indices[r,i]]
        for n in range(Nisom, Nisom+Nreac+Nprod):
            x[s,n] = ode.y[-(Nisom+Nreac+Nprod)+n]

    #import pylab
    #pylab.loglog(t,x)
    #pylab.show()

    return t, p, x

################################################################################

def solveReducedME(T, P, Elist, tlist, x0, K, p0, Nisom, Nreac, Nprod):
    """
    Directly solve a reduced master equation (set of phenomenological rate
    coefficients) using a stiff ODE solver. Pass the reaction `network` to
    solve, the temperature `T` in K and pressure `P` in Pa to solve at, the
    energies `Elist` in J/mol to use, the output time points `tlist` in s,
    the initial total populations `x0`, the matrix of phenomenological rate
    coefficients `K`, and the pseudo-steady population distributions `p0`.
    Returns the times in s, approximate population distributions for each
    isomer, and total population profiles for each configuration.
    """

    Ngrains = len(Elist)
    Ntime = len(tlist)

    # Set up ODEs
    ode = scipy.integrate.ode(residual, jacobian)
    ode.set_integrator('vode', method='bdf', with_jacobian=True, atol=1e-16, rtol=1e-8)
    ode.set_initial_value(x0, 0.0).set_f_params(K).set_jac_params(K)

    # Generate solution
    t = numpy.zeros([len(tlist)], float)
    x = numpy.zeros([len(tlist), len(x0)], float)
    for n in range(Ntime):
        ode.integrate(tlist[n])
        t[n] = ode.t
        x[n,:] = ode.y

    # Construct p from x
    p = numpy.zeros((len(t),Nisom,Ngrains), numpy.float64)
    for s in range(Ntime):
        for i in range(Nisom):
            for n in range(Nisom):
                p[s,i,:] += x[s,n] * p0[:,i,n]
            for n in range(Nisom, Nisom+Nreac):
                p[s,i,:] += x[s,n] * p0[:,i,n] * (0.21 * P / constants.R / T)

    return t, p, x

#################################################################################
#
#class MasterEquation(DASSL):
#    """
#    """
#
#    def __init__(self, M):
#        DASSL.__init__(self)
#        self.M = M
#
#    def residual(self, t, y, dydt):
#        return numpy.dot(self.M, y) - dydt, 0
#
#    def jacobian(self, t, y, dydt, cj):
#        return self.M - cj * numpy.identity(y.shape[0], numpy.float64)