#!/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 sys
import timeit
import hermesv3_gr.config.settings as settings
from warnings import warn as warning


class Speciation(object):
    """
    Speciation class that contains all the needed information to do the speciation.

    :param speciation_id: ID of the speciation profile that have to be in the speciation profile file.
    :type speciation_id: str

    :param speciation_profile_path: Path to the file that contains all the speciation profiles.
    :type speciation_profile_path: str

    :param molecular_weights_path: Path to the file that contains all the needed molecular weights.
    :type molecular_weights_path: str
    """
    def __init__(self, speciation_id, speciation_profile_path, molecular_weights_path):
        st_time = timeit.default_timer()
        settings.write_log('\t\tInitializing Speciation.', level=2)

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

        settings.write_time('Speciation', 'Init', timeit.default_timer() - st_time, level=2)

    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.

        :param speciation_profile_path:
        :type speciation_profile_path:

        :return: List of dictionaries. Each dictionary has the keys 'name', 'formula', 'units' and 'long_name.
        :rtype: list
        """
        import pandas as pd

        st_time = timeit.default_timer()
        settings.write_log("\t\t\tGetting 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:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: Speciation profile ID {0} is not in the {1} file.'.format(
                    self.id, speciation_profile_path))
            sys.exit(1)
        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]
            })

        settings.write_time('Speciation', 'get_speciation_profile', timeit.default_timer() - st_time, level=3)
        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.

        :param molecular_weights_path: Path to the CSV that contains all the molecular weights.
        :type molecular_weights_path: str

        :return: Dictionary with the name of the pollutant as key and the molecular weight as value.
        :rtype: dict
        """
        import pandas as pd

        st_time = timeit.default_timer()

        df = pd.read_csv(molecular_weights_path, sep=',')

        dict_aux = {}

        for i, element in df.iterrows():
            dict_aux.update({element.Specie: element.MW})

        settings.write_time('Speciation', 'extract_molecular_weights', timeit.default_timer() - st_time, level=3)

        return dict_aux

    def do_speciation(self, emission_list):
        """
        Manages all the process to speciate the emissions.

        :param emission_list: List of emissions to speciate.
        :type emission_list: list

        :return: List of emissions already speciated.
        :rtype: list
        """
        import numpy as np

        st_time = timeit.default_timer()
        settings.write_log("\tSpeciating", level=2)

        input_pollutants = []
        # Apply conversion factor to the input pollutants
        for emission in emission_list:
            try:
                emission['data'] = np.array(emission['data'] / self.molecular_weights[emission['name']],
                                            dtype=settings.precision)
            except KeyError:
                settings.write_log('ERROR: Check the .err file to get more info.')
                if settings.rank == 0:
                    raise KeyError('ERROR: {0} pollutant is not in the molecular weights file {1} .'.format(
                        emission['name'], self.molecular_weights_path))
                sys.exit(1)
            exec("{0} = np.array(emission['data'], dtype=settings.precision)".format(emission['name']))
            emission['units'] = ''
            input_pollutants.append(emission['name'])

        del emission_list

        speciated_emissions = []
        num = 0

        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':
                settings.write_log("\t\tPollutant {0} using the formula {0}={3} ({1}/{2})".format(
                    pollutant['name'], num, len(self.speciation_profile), formula), level=3)

                dict_aux = {'name': pollutant['name'],
                            'units': pollutant['units'],
                            'long_name': pollutant['long_name']}
                if formula is '0' or formula is 0:
                    dict_aux.update({'data': 0})
                else:
                    try:
                        dict_aux.update({'data': np.array(eval(formula), dtype=settings.precision)})
                    except NameError as e:
                        settings.write_log('ERROR: Check the .err file to get more info.')
                        if settings.rank == 0:
                            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)))
                        else:
                            sys.exit(1)
                speciated_emissions.append(dict_aux)
            else:
                settings.write_log("\t\tPollutant {0} does not have formula. Ignoring. ({1}/{2})".format(
                    pollutant['name'], num, len(self.speciation_profile)), level=3)
        if len(input_pollutants) > 0:
            settings.write_log("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

    def get_long_name(self, name):

        st_time = timeit.default_timer()
        value = ''
        for pollutant in self.speciation_profile:
            if pollutant['name'] == name:
                value = pollutant['long_name']

        settings.write_time('Speciation', 'get_long_name', timeit.default_timer() - st_time, level=3)
        return value

    def get_units(self, name):

        st_time = timeit.default_timer()
        value = ''

        for pollutant in self.speciation_profile:
            if pollutant['name'] == name:
                value = pollutant['units']

        settings.write_time('Speciation', 'get_units', timeit.default_timer() - st_time, level=3)
        return value


if __name__ == '__main__':
    pass
