# adapted from rcp85.py
#
# Convenience module for loading in RCP (CMIP5) and SSP (CMIP6) emissions datasets
#
# Usage:
#
# import scenario
# scenario.Emissions(ssp="ssp370").co2

import numpy as np
import pandas as pd

import os

# file obtained from: 
# https://www.rcmip.org/
# https://www.google.com/url?q=https%3A%2F%2Frcmip-protocols-au.s3-ap-southeast-2.amazonaws.com%2Fv5.1.0%2Frcmip-emissions-annual-means-v5-1-0.csv&sa=D&sntz=1&usg=AFQjCNHJCbRZjljHC4Y_Qe8QtRlNlXlkNQ
emissions_filename = os.path.join(
    os.path.dirname(__file__), 'data/rcmip-emissions-annual-means-v5-1-0-ssprcp-only.csv')

# TODO: add implementation  of concentrations and forcings


# concentrations_filename = os.path.join(
#     os.path.dirname(__file__), 'data/RCP85_MIDYEAR_CONCENTRATIONS.csv')
# forcing_filename = os.path.join(
#     os.path.dirname(__file__), 'data/RCP85_MIDYEAR_RADFORCING.csv')
# aviNOx_filename = os.path.join(
#     os.path.dirname(__file__), 'data/aviNOx_fraction.csv')
# fossilCH4_filename = os.path.join(
#     os.path.dirname(__file__), 'data/fossilCH4_fraction.csv')

# aviNOx_frac = np.loadtxt(aviNOx_filename, skiprows=5, usecols=(4,),
#     delimiter=',')
# fossilCH4_frac = np.loadtxt(fossilCH4_filename, skiprows=5, usecols=(4,),
#     delimiter=',')

def _convert_emissions_to_fair_units(emissions_df):
    
    from fair.constants import molwt
    
    unit_convert = np.ones(40)
    unit_convert[1] = molwt.C/molwt.CO2/1000   # CO2 fossil
    unit_convert[2] = molwt.C/molwt.CO2/1000   # CO2 land use
    # Notes about N2O
    #   previous (CMIP5) RCPs units: MtN2O-N/yr. 
    #   current (CMIP6) SSP units: ktN20/yr
    unit_convert[4] = molwt.N2/molwt.N2O/1000  # N2O # this is for CMIP6. converts from ktN2O  to T N2
    unit_convert[5] = molwt.S/molwt.SO2  # SO2
    unit_convert[8] = molwt.N/molwt.NO2  # NOx
    
    unit_convert_ser = pd.Series(unit_convert,index=emissions_df.columns)

    emissions_df = emissions_df * unit_convert_ser
    return emissions_df

def _read_emissions(ssp):

    # obtained from https://www.rcmip.org/
    rcemissions = pd.read_csv(emissions_filename)
    emissions = rcemissions.loc[rcemissions.Scenario==ssp]
    emissions = emissions.loc[emissions.Region=='World'].set_index('Variable').drop(
        columns=['Activity_Id', 'Mip_Era', 'Region', 'Model','Scenario', 'Unit']).T
    emissions.index = emissions.index.rename('Years')
    emissions.index = emissions.index.astype('int')
    
    return emissions

def _get_available_scenarios():
    # obtained from https://www.rcmip.org/
    rcemissions = pd.read_csv(emissions_filename)
    return rcemissions.Scenario.unique()

def _get_emissions_units(ssp):
    
    # obtained from https://www.rcmip.org/
    rcemissions = pd.read_csv(emissions_filename)
    rcemissions = rcemissions.loc[(rcemissions.Scenario==ssp)
                               & (rcemissions.Region=='World')].set_index('Variable')['Unit']
    return rcemissions


def _get_ordered_gas_column_names():  
    ''' returns a list of column names from emissions source file, 
    ordered as expected by FaIR emissions array
    '''
    
    # this is the order of the FaIR emissions array
    return [
        'Emissions|CO2|MAGICC Fossil and Industrial',
        'Emissions|CO2|MAGICC AFOLU',
        'Emissions|CH4',
        'Emissions|N2O',
        'Emissions|Sulfur',
        'Emissions|CO',
        'Emissions|VOC',
        'Emissions|NOx', 
        'Emissions|BC', 
        'Emissions|OC',
        'Emissions|NH3',
        'Emissions|F-Gases|PFC|CF4',
        'Emissions|F-Gases|PFC|C2F6',
        'Emissions|F-Gases|PFC|C6F14',
        'Emissions|F-Gases|HFC|HFC23',
        'Emissions|F-Gases|HFC|HFC32',
        'Emissions|F-Gases|HFC|HFC4310mee',
        'Emissions|F-Gases|HFC|HFC125',
        'Emissions|F-Gases|HFC|HFC134a',
        'Emissions|F-Gases|HFC|HFC143a',
        'Emissions|F-Gases|HFC|HFC227ea',
        'Emissions|F-Gases|HFC|HFC245fa',
        'Emissions|F-Gases|SF6',
        'Emissions|Montreal Gases|CFC|CFC11',
        'Emissions|Montreal Gases|CFC|CFC12',
        'Emissions|Montreal Gases|CFC|CFC113',
        'Emissions|Montreal Gases|CFC|CFC114',
        'Emissions|Montreal Gases|CFC|CFC115',
        'Emissions|Montreal Gases|CCl4',
        'Emissions|Montreal Gases|CH3CCl3',
        'Emissions|Montreal Gases|HCFC22',
        'Emissions|Montreal Gases|HCFC141b',
        'Emissions|Montreal Gases|HCFC142b', 
        'Emissions|Montreal Gases|Halon1211',
        'Emissions|Montreal Gases|Halon1202',
        'Emissions|Montreal Gases|Halon1301',
        'Emissions|Montreal Gases|Halon2402',
        'Emissions|Montreal Gases|CH3Br', 
        'Emissions|Montreal Gases|CH3Cl' ]

def _get_ordered_emissions_units(ssp):
    # Returns FaIR emissions units
    
    # rcmip emissions have some gases that require unit conversions
    # so they match what fair expects. Manually update those units to fair units here
    # (see _convert_emissions_to_fair_units() above)
    # Here the first index is a gas (not "year")
    
    rcemissions = _get_emissions_units(ssp=ssp)
    ordered_names = _get_ordered_gas_column_names()
    units = rcemissions[ordered_names]
    units[0] = "Gt C/year"
    units[1] = "Gt C/year"
    units[3] = "tonnes N2/year"
    units[4] = "Mt S/year"
    units[7] = "Mt N/year"
    return units

def _prepare_rcmip_scenario(ssp, interp=True):
    
    emissions = _read_emissions(ssp=ssp)
    
    if "ssp" in ssp:
        # add a methane total column
        emissions['Emissions|CH4'] = (emissions['Emissions|CH4|MAGICC AFOLU']
                                      + emissions['Emissions|CH4|MAGICC Fossil and Industrial'] )

    select_columns_ordered = _get_ordered_gas_column_names()
    
    # select only what we need for FaIR
    fair_emissions = emissions[select_columns_ordered]
    
    # pare down the time series
    fair_emissions = fair_emissions.loc[fair_emissions.index >= 1765]

    # convert emissions units to match what FaIR expects
    fair_emissions = _convert_emissions_to_fair_units(fair_emissions.reset_index())

    # make sure Years are integers
    fair_emissions.Years = fair_emissions.Years.astype(int)

    # interpolate missing emissions
    if interp:
        fair_emissions = fair_emissions.interpolate() 
 
    # check that the shape is as expected
    assert fair_emissions.values.shape == (736,40)

    return fair_emissions

class Emissions:
    
    def __init__(self, scenario):
        
        self.scenario = scenario
        self.emissions = _prepare_rcmip_scenario(ssp=scenario).values
        
        self.year      = self.emissions[:,0]
        self.co2_fossil= self.emissions[:,1]
        self.co2_land  = self.emissions[:,2]
        self.co2       = np.sum(self.emissions[:,1:3],axis=1)
        self.ch4       = self.emissions[:,3]
        self.n2o       = self.emissions[:,4]
        self.sox       = self.emissions[:,5]
        self.co        = self.emissions[:,6]
        self.nmvoc     = self.emissions[:,7]
        self.nox       = self.emissions[:,8]
        self.bc        = self.emissions[:,9]
        self.oc        = self.emissions[:,10]
        self.nh3       = self.emissions[:,11]
        self.cf4       = self.emissions[:,12]
        self.c2f6      = self.emissions[:,13]
        self.c6f14     = self.emissions[:,14]
        self.hfc23     = self.emissions[:,15]
        self.hfc32     = self.emissions[:,16]
        self.hfc43_10  = self.emissions[:,17]
        self.hfc125    = self.emissions[:,18]
        self.hfc134a   = self.emissions[:,19]
        self.hfc143a   = self.emissions[:,20]
        self.hfc227ea  = self.emissions[:,21]
        self.hfc245fa  = self.emissions[:,22]
        self.sf6       = self.emissions[:,23]
        self.cfc11     = self.emissions[:,24]
        self.cfc12     = self.emissions[:,25]
        self.cfc113    = self.emissions[:,26]
        self.cfc114    = self.emissions[:,27]
        self.cfc115    = self.emissions[:,28]
        self.carb_tet  = self.emissions[:,29]
        self.mcf       = self.emissions[:,30]
        self.hcfc22    = self.emissions[:,31]
        self.hcfc141b  = self.emissions[:,32]
        self.hcfc142b  = self.emissions[:,33]
        self.halon1211 = self.emissions[:,34]
        self.halon1202 = self.emissions[:,35]
        self.halon1301 = self.emissions[:,36]
        self.halon2402 = self.emissions[:,37]
        self.ch3br     = self.emissions[:,38]
        self.ch3cl     = self.emissions[:,39]
    
        self.gas_species_names = ['CO2_Fossil',
                'CO2_Land',
                'CH4',
                'N2O',
                'SOx',
                'CO',
                'NMVOC',
                'NOx', 
                'BC', 
                'OC',
                'NH3',
                'CF4',
                'C2F6',
                'C6F14',
                'HFC23',
                'HFC32',
                'HFC4310mee',
                'HFC125',
                'HFC134a',
                'HFC143a',
                'HFC227ea',
                'HFC245fa',
                'SF6',
                'CFC11',
                'CFC12',
                'CFC113',
                'CFC114',
                'CFC115',
                'CARB_TET',
                'MCF',
                'HCFC22',
                'HCFC141b',
                'HCFC142b', 
                'Halon1211',
                'Halon1202',
                'Halon1301',
                'Halon2402',
                'CH3Br', 
                'CH3Cl' ]
    
    def get_gas_species_names(self):
        
        return self.gas_species_names
    
    
    def get_gas_species_units(self):
        # return the gas unit strings
        units = _get_ordered_emissions_units(ssp=self.scenario).values
        
        return list(units)
    
    def get_emissions_unit_labels(self):
        # returns unit label string for all 40 members of FaIR emissions (including year at idx=0)
        gas_names = self.get_gas_species_names()
        
        return ['Year'] + gas_names
    
    def get_emissions_array_labels(self):
        
        gas_names = self.get_gas_species_names()
        
        return ['Year'] + gas_names
    
    def get_available_scenarios():
        
        return _get_available_scenarios()

class Concentrations:
    
    gas_species_names = ['co2',
              'ch4',
              'n2o',
              'cf4',
              'c2f6',
              'c6f14',
              'hfc23',
              'hfc32',
              'hfc43_10',
              'hfc125',
              'hfc134a',
              'hfc143a',
              'hfc227ea',
              'hfc245fa',
              'sf6',
              'cfc11',
              'cfc12',
              'cfc113',
              'cfc114',
              'cfc115',
              'carb_tet',
              'mcf',
              'hcfc22',
              'hcfc141b',
              'hcfc142b',
              'halon1211',
              'halon1202',
              'halon1301',
              'halon2402',
              'ch3br',
              'ch3cl'
             ]

    def get_concentrations_array_labels():
        
        return Concentrations.get_gas_species_names()

    def get_gas_species_names():
        return Concentrations.gas_species_names


#     concentrations = np.loadtxt(concentrations_filename, skiprows=39,
#       delimiter=',')
#     gas_indices= np.concatenate(([3,4,5], np.arange(8,36)))
#     gases      = concentrations[:,gas_indices]
#     year       = concentrations[:,0]
#     co2eq      = concentrations[:,1]
#     kyotoco2eq = concentrations[:,2]
#     co2        = concentrations[:,3]
#     ch4        = concentrations[:,4]
#     n2o        = concentrations[:,5]
#     fgassum    = concentrations[:,6]
#     mhalosum   = concentrations[:,7]
#     cf4        = concentrations[:,8]
#     c2f6       = concentrations[:,9]
#     c6f14      = concentrations[:,10]
#     hfc23      = concentrations[:,11]
#     hfc32      = concentrations[:,12]
#     hfc43_10   = concentrations[:,13]
#     hfc125     = concentrations[:,14]
#     hfc134a    = concentrations[:,15]
#     hfc143a    = concentrations[:,16]
#     hfc227ea   = concentrations[:,17]
#     hfc245fa   = concentrations[:,18]
#     sf6        = concentrations[:,19]
#     cfc11      = concentrations[:,20]
#     cfc12      = concentrations[:,21]
#     cfc113     = concentrations[:,22]
#     cfc114     = concentrations[:,23]
#     cfc115     = concentrations[:,24]
#     carb_tet   = concentrations[:,25]
#     mcf        = concentrations[:,26]
#     hcfc22     = concentrations[:,27]
#     hcfc141b   = concentrations[:,28]
#     hcfc142b   = concentrations[:,29]
#     halon1211  = concentrations[:,30]
#     halon1202  = concentrations[:,31]
#     halon1301  = concentrations[:,32]
#     halon2402  = concentrations[:,33]
#     ch3br      = concentrations[:,34]
#     ch3cl      = concentrations[:,35]    


class Forcing:

    forcing_names = ['co2',
                 'ch4',
                 'n2o',
                 'fgas',
                 'halo',
                 'aero',
                 'cloud',
                 'strato3',
                 'tropo3',
                 'stwv',
                 'dust',
                 'landuse',
                 'bcsnow'
                ]

    def get_forcing_names():
        return Forcing.forcing_names
    
    def get_forcing_array_labels():
        
        return Forcing.get_forcing_names()

#     forcing   = np.loadtxt(forcing_filename, skiprows=59, delimiter=',')
#     year      = forcing[:,0]
#     total     = forcing[:,1]
#     volcanic  = forcing[:,2]
#     solar     = forcing[:,3]
#     ghg       = forcing[:,5]
#     co2       = forcing[:,8]
#     ch4       = forcing[:,9]
#     n2o       = forcing[:,10]
#     fgas      = forcing[:,11]
#     halo      = forcing[:,12]
#     aero      = forcing[:,41]
#     cloud     = forcing[:,48]
#     strato3   = forcing[:,49]
#     tropo3    = forcing[:,50]
#     stwv      = forcing[:,51]
#     dust      = forcing[:,47]
#     landuse   = forcing[:,52]
#     bcsnow    = forcing[:,53]

