#!/usr/bin/env python

# Copyright 2018 Earth Sciences Department, BSC-CNS
#
# This file is part of HERMESv3_GR.
#
# HERMESv3_GR is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# HERMESv3_GR is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with HERMESv3_GR. If not, see <http://www.gnu.org/licenses/>.


import timeit
import pandas as pd
import numpy as np
from hermesv3_gr.config import settings, log_message
from nes import Nes
from warnings import warn as warning


class Speciation(object):
    """
    Attributes
    ----------
    speciation_profile : List[dict]
        List of dictionaries. Each dictionary has the keys 'name', 'formula', 'units' and 'long_name'.
    molecular_weights : dict
        Dictionary with the name of the pollutant as key and the molecular weight as value.
    """
    def __init__(self, speciation_id, speciation_profile_path, molecular_weights_path):
        """
        Speciation class initialization

        Parameters
        ----------
        speciation_id : str
            ID of the speciation profile that have to be in the speciation profile file.
        speciation_profile_path : str
            Path to the file that contains all the speciation profiles.
        molecular_weights_path : str
            Path to the file that contains all the needed molecular weights.
        """

        self.id = speciation_id
        self.speciation_profile = self.get_speciation_profile(speciation_profile_path)
        self.molecular_weights = self.extract_molecular_weights(molecular_weights_path)

    def get_speciation_profile(self, speciation_profile_path):
        """
        Extract the speciation information as a dictionary with the destiny pollutant as key and the formula as value.

        Parameters
        ----------
        speciation_profile_path : str
            Path to the speciation profile file.

        Returns
        -------
        List[dict]
            List of dictionaries. Each dictionary has the keys 'name', 'formula', 'units' and 'long_name'.
        """
        log_message("Getting speciation profile id '{0}' from {1} .".format(
            self.id, speciation_profile_path), level=3)

        df = pd.read_csv(speciation_profile_path, sep=';')

        try:
            formulas_dict = df.loc[df[df.ID == self.id].index[0]].to_dict()
        except IndexError:
            raise AttributeError('ERROR: Speciation profile ID {0} is not in the {1} file.'.format(
                self.id, speciation_profile_path))
        formulas_dict.pop('ID', None)
        units_dict = df.loc[df[df.ID == 'units'].index[0]].to_dict()

        units_dict.pop('ID', None)
        long_name_dict = df.loc[df[df.ID == 'short_description'].index[0]].to_dict()
        long_name_dict.pop('ID', None)
        profile_list = []
        for key in formulas_dict.keys():
            profile_list.append({
                'name': key,
                'formula': formulas_dict[key],
                'units': units_dict[key],
                'long_name': long_name_dict[key]
            })

        return profile_list

    @staticmethod
    def extract_molecular_weights(molecular_weights_path):
        """
        Extract the molecular weights for each pollutant as a dictionary with the name of the pollutant as key and the
        molecular weight as value.

        Parameters
        ----------
        molecular_weights_path : str
            Path to the CSV that contains all the molecular weights.
        Returns
        -------
        dict
            Dictionary with the name of the pollutant as key and the molecular weight as value.
        """

        df = pd.read_csv(molecular_weights_path, sep=',', index_col='Specie',
                         dtype={'Specie': str, 'MW': settings.precision})
        molecular_weights = df.loc[:, 'MW']

        return molecular_weights

    def do_speciation(self, emissions):
        """
        Speciate the emissions

        Parameters
        ----------
        emissions : Nes
            Emissions NES object

        Returns
        -------
        Nes
            Speciated emissions NES object
        """
        st_time = timeit.default_timer()
        log_message("Speciating", level=2)

        for var_name in emissions.variables.keys():
            try:
                emissions.variables[var_name]['data'] = np.array(
                    emissions.variables[var_name]['data'] / self.molecular_weights[var_name], dtype=settings.precision)
            except KeyError:
                raise KeyError('ERROR: {0} pollutant is not in the molecular weights file.'.format(var_name))
            exec("{0} = np.array(emissions.variables[var_name]['data'], dtype=settings.precision)".format(var_name))

        speciated_emissions = emissions.copy(copy_vars=False)
        num = 0
        input_pollutants = list(emissions.variables.keys())
        for pollutant in self.speciation_profile:
            formula = str(pollutant['formula'])
            used_poll = []
            for in_p in input_pollutants:
                if in_p in formula:
                    used_poll.append(in_p)
            for poll_rem in used_poll:
                input_pollutants.remove(poll_rem)
            num += 1
            if formula != 'nan':
                log_message("Pollutant {0} using the formula {0}={3} ({1}/{2})".format(
                    pollutant['name'], num, len(self.speciation_profile), formula), level=3)

                dict_aux = {'units': pollutant['units'],
                            'long_name': pollutant['long_name']}
                if formula == '0' or formula == 0:
                    dict_aux.update({'data': 0})
                else:
                    try:
                        dict_aux.update({'data': np.array(eval(formula), dtype=settings.precision)})
                    except NameError as e:
                        raise AttributeError(
                            "Error in speciation profile {0}: ".format(self.id) +
                            "The output specie {0} cannot be calculated ".format(pollutant['name']) +
                            "with the expression {0} because{1}".format(formula, str(e)))

                speciated_emissions.variables[pollutant['name']] = dict_aux
            else:
                log_message("Pollutant {0} does not have formula. Ignoring. ({1}/{2})".format(
                    pollutant['name'], num, len(self.speciation_profile)), level=3)
        if len(input_pollutants) > 0:
            log_message("WARNING: The input pollutants {0} do not appear in the speciation profile {1}.".format(
                    input_pollutants, self.id))
            if settings.rank == 0:
                warning("WARNING: The input pollutants {0} do not appear in the speciation profile {1}.".format(
                    input_pollutants, self.id))
        settings.write_time('Speciation', 'do_speciation', timeit.default_timer() - st_time, level=3)

        return speciated_emissions
