"""
Generic PCM Photochemistry post-processing library
Written by Maxime Maurice in 2024 anno domini
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as tk
import xarray as xr
from scipy.constants import R, N_A

import warnings
warnings.filterwarnings("ignore", message="The following kwargs were not used by contour: 't'")
warnings.filterwarnings("ignore", message="The following kwargs were not used by contour: 'lon'")
warnings.filterwarnings("ignore", message="The following kwargs were not used by contour: 'lat'")

M          = {'co2':44,  # Molar masses
              'o':16,    # TODO: automatic parser
              'o1d':16,
              'o2':32,
              'o3':48,
              'h':1,
              'h2':2,
              'oh':17,
              'h2o_vap':18,
              'ho2':33,
              'h2o2':34,
              'co':28,
              'cho':29,
              'ch2o':30}

background = 'co2' # background gas of the atmosphere

class GPCM_simu:
    """ Generic PCM Simulation

    Stores the netCDF file and path to the simulation

    Attributes
    ----------
    data : xr.Dataset
        NetCDF file of the simulation
    path : str
        Path to simulation directory

    Methods
    -------
    get_profile(field,**kw)
        Get profile of a field (either local or averaged)
    plot_meridional_slice(field,logcb,labelcb,**kw)
        Plot a meridional slice of a field
    plot_time_evolution(field,logcb,labelcb,**kw)
        Plot time evolution of a field (as a time vs altitude contour plot)
    plot_profile(field,**kw)
        Plot a profile of a field

    """
    def __init__(self,path,datafilename='diagfi',verbose=False):
        """
        Parameters
        ----------
        path : str
            Path to simulation directory
        datafilename : str (optional)
            Name of the netCDF file (by default: diagfi)
            works with start, startfi, concat etc.
            Do not add .nc type suffix
        """
        self.path = path
        try:
            self.data = xr.open_dataset(path+'/'+datafilename+'.nc',decode_times=False)
            print(path+'/'+datafilename,'loaded, simulations lasts',self.data['Time'].values[-1],'sols')
        except:
            raise Exception('Data not found')

    def __getitem__(self,field):
        return self.data[field]

    def __setitem__(self,field,value):
        self.data[field] = value

    def get_profile(self,field,**kw):
        """ Get profile of a field (either local or averaged)

        Parameters
        ----------
        field : str
            Field name
        t : float (optional)
            Time at which to select (if nothing specified use time-average)
        lat : float (optional)
            Latitude at which to select (if nothing specified use area-weighted meridional average)
        lon : float (optional)
            Longitude at which to select (if nothing specified use zonal average)

        """
        
        if self['latitude'].size == 1:
            # 1D
            return self[field][:,0,0]
        else:
            # 3D
            if 'lat' in kw:
                if 'lon' in kw:
                    return self[field].sel(latitude=kw['lat'],method='nearest').sel(longitude=kw['lon'],method='nearest')
                else:
                    return self[field].sel(latitude=kw['lat'],method='nearest').mean(dim='longitude')
            else:
                # Latitude-averaged profile: need to weight by the grid cell surface area
                if 'lon' in kw:
                    return (self['aire']*self[field]/self['aire'].mean(dim='latitude')).mean(dim='latitude').sel(longitude=kw['lon'],method='nearest')
                else:
                    return (self['aire']*self[field]/self['aire'].mean(dim='latitude').mean(dim='longitude')).mean(dim='latitude').mean(dim='longitude')

    def get_subset(self,field='all',**kw):
        """ Get a subset at fixed given coordinate of the data

        Parameters
        ----------
        field : str (optional, default = 'all')
            Field name. If nothing or 'all'
            specified, return all fields.
        t : float (optional)
            Time of the slice. If nothing
            specified, use time average
        lon : float (optional)
            Longitude of the slice. If nothing
            specified, use meridional average
        lat : float (optional)
            Latitude of the slice. If nothing
            specified, use area-weighted zonal average
        alt : float (optional)
            Altitude of the slice. If nothing
            specified, use time-average

        Raise
        -----
        Slice direction not provided
        """
        if len(kw) == 0:
            raise Exception('Slice direction not provided')

        if field == 'all':
            data_subset = self.data
        else:
            data_subset = self[field]
            
        if 't' in kw:
            data_subset = data_subset.sel(Time=kw['t'],method='nearest')
        if 'lon' in kw:
            data_subset = data_subset.sel(longitude=kw['lon'],method='nearest')
        if 'lat' in kw:
            data_subset = data_subset.sel(latitude=kw['lat'],method='nearest')
        if 'alt' in kw:
            data_subset = data_subset.sel(altitude=kw['alt'],method='nearest')

        return data_subset

    def plot_meridional_slice(self,field,t='avg',lon='avg',logcb=False,labelcb=None,**plt_kw):
        """ Plot a meridional slice of a field

        Parameters
        ----------
        field : str
            Field name to plot
        logcb : bool (optional)
            Use logarithmic colorscale
        labelcb : str (optional)
            Use custom colorbar label
        t : float (keyword)
            Time at which to plot (if nothing specified use time average)
        lon : float (keyword)
            Longitude at which to plot (if nothing specified use zonal average)
        """

        if self['latitude'].size == 1:
            # safety check
            raise Exception('Trying to plot a meridional slice of a 1D simulation')

        meridional_slice = self[field]
        
        if t == 'avg':
            meridional_slice = meridional_slice.mean(dim='Time')
        else:
            meridional_slice = meridional_slice.sel(Time=t,method='nearest')
        if lon == 'avg':
            meridional_slice = meridional_slice.mean(dim='longitude')
        else:
           meridional_slice = meridional_slice.sel(longitude=lon,method='nearest') 
            
        if logcb:
            plt.contourf(self['latitude'],self['altitude'],meridional_slice,locator=tk.LogLocator(),**plt_kw)
        else:
            plt.contourf(self['latitude'],self['altitude'],meridional_slice,**plt_kw)
            
        plt.colorbar(label=field if labelcb==None else labelcb)
        plt.xlabel('latitude [°]')
        plt.ylabel('altitude [km]')

    def plot_time_evolution(self,field,lat='avg',lon='avg',logcb=False,labelcb=None,**plt_kw):
        """ Plot time evolution of a field (as a time vs altitude contour plot)

        Parameters
        ----------
        field : str
            Field name to plot
        lat : float (optional)
            Latitude at which to plot (if nothing specified use area-weighted meridional average)
        lon : float (optional)
            Longitude at which to plot (if nothing specified use zonal average)
        logcb : bool (optional)
            Use logarithmic colorscale
        labelcb : str (optional)
            Use custom colorbar label
        matplotlib contourf keyword arguments
        """

        time_evolution = self[field]
        
        if lat == 'avg':
            time_evolution = (self['aire']*time_evolution/self['aire'].mean(dim='latitude')).mean(dim='latitude')
        else:
            time_evolution = time_evolution.sel(latitude=lat,method='nearest')
        if lon == 'avg':
            time_evolution = time_evolution.mean(dim='longitude')
        else:
            time_evolution = time_evolution.sel(longitude=lon,method='nearest')
            
        if logcb:
            plt.contourf(self['Time'],self['altitude'],time_evolution.T,locator=tk.LogLocator(),**plt_kw)
        else:
            plt.contourf(self['Time'],self['altitude'],time_evolution.T,**plt_kw)
            
        plt.colorbar(label=field if labelcb==None else labelcb)
        plt.xlabel('time [day]')
        plt.ylabel('altitude [km]')

    def plot_atlas(self,field,t='avg',alt='avg',logcb=False,labelcb=None,**plt_kw):
        """ Plot atlas of a field

        Parameters
        ----------
        field : str
            Field name to plot
        t : float (optional)
            Time at which to pot (if nothing specified, use time average)
        alt : float (optional)
            Altitude at which to plot (if nothing specified use vertical average)
        logcb : bool (optional)
            Use logarithmic colorscale
        labelcb : str (optional)
            Use custom colorbar label
        matplotlib contourf keyword arguments
        """

        atlas = self[field]
        
        if t == 'avg':
            atlas = atlas.mean(dim='Time')
        else:
            atlas = atlas.sel(Time=t,method='nearest')
            
        if  'altitude' in atlas.coords:
            if alt == 'avg':
                atlas = atlas.mean(dim='altitude')
            else:
                atlas = atlas.sel(altitude=alt,method='nearest')
            
        if logcb:
            plt.contourf(self['longitude'],self['latitude'],atlas,locator=tk.LogLocator(),**plt_kw)
        else:
            plt.contourf(self['longitude'],self['latitude'],atlas,**plt_kw)
            
        plt.colorbar(label=field if labelcb==None else labelcb)
        plt.xlabel('longitude [°]')
        plt.ylabel('matitude [°]')

    def plot_profile(self,field,t='avg',lon='avg',lat='avg',logx=False,**plt_kw):
        """ Plot a profile of a field

        Parameters
        ----------
        field : str
            Field name to plot
        logx : bool (optional)
            Use logarithmic x axis
        t : float (optional)
            Time at which to select (if nothing specified use time-average)
        lat : float (optional)
            Latitude at which to plot (if nothing specified use area-weighted meridional average)
        lon : float (optional)
            Longitude at which to plot (if nothing specified use zonal average)
        matplotlib's plot / semilogx keyword arguments
        """

        profile = self[field]
        
        if t == 'avg':
            profile = profile.mean(dim='Time')
        else:
            profile = profile.sel(Time=t,method='nearest')
            
        if self['latitude'].size > 1:
            
            if lat == 'avg':
                profile = (self['aire']*profile/self['aire'].mean(dim='latitude')).mean(dim='latitude')
            else:
                profile = profile.sel(latitude=lat,method='nearest')
            if lon == 'avg':
                profile = profile.mean(dim='longitude')
            else:
                profile = profile.sel(longitude=lon,method='nearest')
            
        if logx:
            plt.semilogx(profile,self['altitude'],**plt_kw)
        else:
            plt.plot(profile,self['altitude'],**plt_kw)
            
        plt.xlabel(field+' ['+self[field].units+']')
        plt.ylabel('altitude [km]')
            
class reaction:
    """ Instantiates a basic two-body reaction
    
    Attributes
    ----------
    formula : str
        Reaction formula (e.g. "A + B -> C + D")
    reactants : list
        Reactanting molecules formulae (e.g. ["A", "B"])
    products : list
        Produced molecules formulae (e.g. ["C", "D"])
    constant : fun
        Reaction rate constant, function of potentially temperature and densities

    Methods
    -------
    rate(T,densities)
        Reaction rate for given temperature and densities
    """
    
    def __init__(self,reactants,products,constant):
        """
        Parameters
        ----------
        reactants : list
            Reacting molecules formulae (e.g. ["A", "B"])
        products : list
            Produced molecules formulae (e.g. ["C", "D"])
        constant : fun
            Reaction rate constant, function of potentially temperature and densities
        """
        
        self.formula   = ''.join([r_+' + ' for r_ in reactants[:-1]])+reactants[-1]+' -> '+''.join([r_+' + ' for r_ in products[:-1]])+products[-1]
        self.products  = products
        self.reactants = reactants
        self.constant  = constant

    def rate(self,T,densities):
        """ Computes reaction rate

        Parameters
        ----------
        T : float
            Temperature [K]
        densities : dict
            Molecular densities [cm^-3]

        Returns
        -------
        float
            Value of the reaction rate [cm^-3.s^-1]
        """
        
        return self.constant(T,densities[background])*densities[self.reactants[0]]*densities[self.reactants[1]]

class termolecular_reaction(reaction):

    def __init__(self,reactants,products,constant):

        self.formula   = ''.join([r_+' + ' for r_ in reactants[:-1]])+reactants[-1]+' + M -> '+''.join([r_+' + ' for r_ in products[:-1]])+products[-1]+' + M'
        self.products  = products
        self.reactants = reactants
        self.constant  = constant

class photolysis(reaction):

    def __init__(self,reactants,products):

        self.formula   = ''.join([r_+' + ' for r_ in reactants[:-1]])+reactants[-1]+' + hv -> '+''.join([r_+' + ' for r_ in products[:-1]])+products[-1]
        self.products  = products
        self.reactants = reactants

    def rate(self,j,densities):
        """ Computes reaction rate

        Parameters
        ----------
        j : float
            Photolysis rate [s^-1]
        densities : dict
            Molecular densities [cm^-3]

        Returns
        -------
        float
            Value of the reaction rate [cm^-3.s^-1]
        """

        return j*densities[self.reactants[0]]

class reaction_constant:
    """ Basic (Arrhenius) reaction rate constant
    
    Instantiates type 1 rate constant for a particular reaction
    (https://lmdz-forge.lmd.jussieu.fr/mediawiki/Planets/index.php/Photochemistry#Reaction_rate_formulae)
    
    Attributes
    ----------
    params : dict
        Reaction-specific set of parameters for the rate constant: a,T0,c,b,d

    Methods
    -------
    call(T,density)
        Compute the reaction rate for given temperature and density
    """

    def __init__(self,params):
        """
        Parameters
        ----------
        params : dict
            Reaction-specific set of parameters for the rate constant: a,T0,c,b,d
        """
        self.params = params

    def __call__(self,T,density):
        """
        Parameters
        ----------
        T : float
            Temperature [K]
        density : float
            Background gas density [cm^-3]

        Returns
        -------
        float
            Value of the reaction rate constant [cm^3.s^-1]
        """
        return self.params['a']*(T/self.params['T0'])**self.params['c']*np.exp(-self.params['b']/T)*density**self.params['d']

class reaction_constant_type2(reaction_constant):
    """ Type 2 reaction rate constant
    
    Instantiates type 2 rate constant for a particular reaction
    (https://lmdz-forge.lmd.jussieu.fr/mediawiki/Planets/index.php/Photochemistry#Reaction_rate_formulae)
    
    Attributes
    ----------
    params : dict
        Reaction-specific set of parameters for the rate constant: k0,T0,n,a0,kinf,m,b0,g,h,dup,ddown,fc

    Methods
    -------
    call(T,density)
        Computes the reaction rate for given temperature and density
    """
    def __call__(self,T,density):
        """
        Parameters
        ----------
        T : float
            Temperature [K]
        density : float
            Background gas density [cm^-3]

        Returns
        -------
        float
            The value of the reaction rate constant [cm^3.s^-1]
        """
        num = self.params['k0']*(T/self.params['T0'])**self.params['n']*np.exp(-self.params['a0']/T)
        den = self.params['kinf']*(T/self.params['T0'])**self.params['m']*np.exp(-self.params['b0']/T)
    
        return self.params['g']*np.exp(-self.params['h']/T)+num*density**self.params['dup']/(1+num/den*density**self.params['ddown'])*self.params['fc']**(1/(1+np.log10(num/den*density)**2))

# TODO: implement type 3 reaction constant

def read_reactfile(path):
    """ Reads the reactfile formatted for simulations with the Generic PCM
    
    Parameters
    ----------
    path : str
        Path to the reactfile (to become reaction.def)

    Returns
    -------
    dict
        Keys are reactions formulae, items are reactions instances
    """

    reactions = {}
    with open(path+'/chemnetwork/reactfile') as reactfile:
        for line in reactfile:
            # Commented line
            if line[0] == '!':
                # Hard-coded reaction
                if 'hard' in line and 'coded' in line:
                    hard_coded_reaction = reaction(line[1:51].split(),line[51:101].split(),None)
                    print('reaction ',hard_coded_reaction.formula,'seems to be hard-coded. Add it manually if needed.')
                continue
            else:
                reactants = line[:50].split()
                products  = line[50:100].split()
                
                # Photolysis
                if 'hv' in line:
                    reactants.pop(reactants.index('hv'))
                    new_reaction = photolysis(reactants,products)
                    
                # Other reactions
                else:
                    cst_params = [float(val) for val in line[101:].split()]
                    
                    # three-body reaction
                    if 'M' in line:
                        
                        # if third body is not the background gas
                        if line[line.index('M')+2] != ' ':
                            third_body = reactants[reactants.index('M')+1]
                        else:
                            third_body = 'background'
                        reactants.pop(reactants.index('M'))
                        try:
                            products.pop(reactants.index('M'))
                        except:
                            pass
                        if int(line[100]) == 1:
                            rate_constant = reaction_constant({param_key:cst_params[i] for i,param_key in enumerate(['a','b','c','T0','d'])})
                            # if the third body is not the background gas, we treat it as a two-body reaction
                            if third_body != 'background':
                                products.append(third_body)
                                rate_constant.params['d'] = 0
                        elif int(line[100]) == 2:
                            rate_constant = reaction_constant_type2({param_key:cst_params[i] for i,param_key in enumerate(['k0','n','a0','kinf','m','b0','T0','fc','g','h','dup','ddown'])})
                            if third_body != 'background':
                                raise Exception('Dont know how to handle three body reaction with type 2 constant when third body is not the background gas.')
                        else:
                            raise Exception('rate constant parameterization type ',line[100],' not recognized.')
                        new_reaction = termolecular_reaction(reactants,products,rate_constant)

                    # two-body reaction
                    else:
                        rate_constant = reaction_constant({param_key:cst_params[i] for i,param_key in enumerate(['a','b','c','T0','d'])})
                        new_reaction = reaction(reactants,products,rate_constant)
                
                reactions[new_reaction.formula] = new_reaction

    return reactions
                

def density(P,T,VMR=1):
    """ Computes molecular density using the perfect gas law

    Parameters
    ----------
    P : float
        Pressure [Pa]
    T : float
        Temperature [K]
    VMR : float (optional)
        Volume mixing ratio [cm^3/cm^3]

    Returns
    -------
    float
        Molecular density [cm^-3]
    """
    
    m3_to_cm3 = 1e6
    return VMR * P / R / T / m3_to_cm3 * N_A

def compute_rates(s,reactions='read'):
    """ Computes reaction rates for a simulation

    Parameters
    ----------
    s : GPCM_simu
        Simulation object
    reactions : dict or 'read' (optional)
        Dictionnary of reactions whose rate to compute as returned by read_reactfile
        If nothing passed, call read_reactfile to identify reactions

    Returns
    -------
    GPCM_simu
        Simulation object with reactions rates, rates constants, species vmr and densities as new fields
    """

    if reactions == 'read':
        reactions      = read_reactfile(s.path)
        s.species   = []
        s.reactions = {} # reactions dict will be merged at the end

    # Register new species
    for r in reactions:
        for sp in reactions[r].reactants:
            if not sp in s.species:
                s.species.append(sp)
        for sp in reactions[r].products:
            if not sp in s.species:
                s.species.append(sp)
        
    densities = {}

    # Background density
    s['total density']     = density(s['p'],s['temp']).assign_attrs({'units':'cm^-3.s^-1'})

    for sp in s.species:
        # volume mixing ratios
        s[sp+' vmr'] = s[sp] * M['co2'] / M[sp]
        s[sp+' vmr'] = s[sp+' vmr'].assign_attrs({'units':'m^3/m^3'})
        # molecular densities
        s[sp+' density'] = density(s['p'],s['temp'],VMR=s[sp+' vmr'])
        s[sp+' density'] = s[sp+' density'].assign_attrs({'units':'cm^-3'})
        densities[sp] = s[sp+' density']

    for r in reactions:

        # Photolysis
        if type(reactions[r]) == photolysis:

            # Cases with branching ratios
            if reactions[r].reactants[0] == 'co2':
                if 'o1d' in reactions[r].products:
                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jco2_o1d'],densities)
                else:
                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jco2_o'],densities)
            elif reactions[r].reactants[0] == 'o2':
                if 'o1d' in reactions[r].products:
                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jo2_o1d'],densities)
                else:
                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jo2_o'],densities)
            elif reactions[r].reactants[0] == 'o3':
                if 'o1d' in reactions[r].products:
                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jo3_o1d'],densities)
                else:
                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jo3_o'],densities)
            elif reactions[r].reactants[0] == 'ch2o':
                if 'cho' in reactions[r].products:
                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jch2o_cho'],densities)
                else:
                    s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jch2o_co'],densities)
            elif reactions[r].reactants[0] == 'h2o_vap':
                s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['jh2o'],densities)
            else:
                # General case
                s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['j'+reactions[r].reactants[0]],densities)
        else:
            s['k ('+reactions[r].formula+')'] = reactions[r].constant(s['temp'],densities[background])
            s['rate ('+reactions[r].formula+')'] = reactions[r].rate(s['temp'],densities)

            # Termolecular reaction
            if type(reactions[r]) == termolecular_reaction:
                s['k ('+reactions[r].formula+')'] = s['k ('+reactions[r].formula+')'].assign_attrs({'units':'cm^6.s^-1'})

            # Bimolecular reaction
            else:
                s['k ('+reactions[r].formula+')'] = s['k ('+reactions[r].formula+')'].assign_attrs({'units':'cm^3.s^-1'})
                
        s['rate ('+reactions[r].formula+')'] = s['rate ('+reactions[r].formula+')'].assign_attrs({'units':'cm^-3.s^-1'})
                
    s.reactions = s.reactions | reactions
    
    return s