#!/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 os
import sys
import timeit

import hermesv3_gr.config.settings as settings
from hermesv3_gr.modules.regrid.regrid_conservative import ConservativeRegrid
from hermesv3_gr.modules.vertical.vertical import VerticalDistribution
from hermesv3_gr.modules.temporal.temporal import TemporalDistribution
from hermesv3_gr.modules.speciation.speciation import Speciation


class EmissionInventory(object):
    """
    Class that defines the content and the methodology for the area emission inventories

    :param current_date: Date to simulate.
    :type current_date: datetime.datetime

    :param inventory_name: Name of the inventory to use.
    :type inventory_name: str

    :param sector: Name of the sector of the inventory to use.
    :type sector: str

    :param pollutants: List of the pollutant name to take into account.
    :type pollutants: list of str

    :param inputs_path: Path where are stored all the datasets to use. eg: /esarchive/recon/jrc/htapv2/monthly_mean
    :type inputs_path: str

    :param input_frequency: Frequency of the inputs. [yearly, monthly, daily]
    :type input_frequency: str

    :param reference_year: year of reference of the information of the dataset.
    :type reference_year: int

    :param factors: Description of the scale factors per country. (e.g. SPN 1.5, CHN 3.)
    :type factors: str

    :param regrid_mask: Description of the masking countries (adding e.g. + SPN AND) (subtracting e.g. - SPN)
    :type regrid_mask: str

    :param p_vertical: ID of the vertical profile to use.
    :type p_vertical: str

    :param p_month: ID of the temporal monthly profile to use.
    :type p_month: str

    :param p_week: ID of the temporal daily profile to use.
    :type p_week: str

    :param p_hour: ID of the temporal hourly profile to use.
    :type p_hour: str

    :param p_speciation: ID of the speciation profile to use.
    :type p_speciation: str
    """
    def __init__(self, options, grid, current_date, inventory_name, source_type, sector, pollutants, inputs_path,
                 input_frequency, vertical_output_profile, reference_year=2010, factors=None, regrid_mask=None,
                 p_vertical=None, p_month=None, p_week=None, p_day=None, p_hour=None, p_speciation=None):
        from hermesv3_gr.modules.masking.masking import Masking

        st_time = timeit.default_timer()
        settings.write_log('\t\tCreating area source emission inventory.', level=3)

        # Emission Inventory parameters
        self.source_type = source_type
        self.date = current_date
        self.inventory_name = inventory_name
        self.sector = sector
        self.reference_year = reference_year
        self.inputs_path = inputs_path
        self.input_frequency = input_frequency
        self.grid = grid

        # Profiles
        p_vertical = self.get_profile(p_vertical)
        p_month = self.get_profile(p_month)
        p_week = self.get_profile(p_week)
        p_hour = self.get_profile(p_hour)
        p_speciation = self.get_profile(p_speciation)

        # Creating Masking Object
        # It will also create the WoldMasks necessaries
        self.masking = Masking(
            options.world_info, factors, regrid_mask, grid,
            world_mask_file=os.path.join(os.path.dirname(options.auxiliar_files_path),
                                         '{0}_WorldMask.nc'.format(inventory_name)))

        self.input_pollutants = pollutants
        self.pollutant_dicts = self.create_pollutants_dicts(pollutants)

        self.masking.check_regrid_mask(self.pollutant_dicts[0]['path'])

        # Creating Regrid Object
        # It will also create the WoldMasks necessaries
        if self.source_type == 'area':
            self.regrid = ConservativeRegrid(
                self.pollutant_dicts,
                os.path.join(options.auxiliar_files_path,
                             "Weight_Matrix_{0}_{1}.nc".format(self.inventory_name, settings.size)),
                grid, masking=self.masking)

        # Creating Vertical Object
        if p_vertical is not None:
            self.vertical = VerticalDistribution(
                self.get_profile(p_vertical), vertical_profile_path=options.p_vertical,
                vertical_output_profile=vertical_output_profile)
        else:
            self.vertical = None
            settings.write_log('\t\tNone vertical profile set.', level=2)
        self.vertical_factors = None

        # Creating Temporal Object
        # It will also create the necessaries timezone files
        if not((p_month is None) and (p_week is None) and (p_day is None) and (p_hour is None)):
            self.temporal = TemporalDistribution(
                current_date, options.output_timestep_type, options.output_timestep_num, options.output_timestep_freq,
                options.p_month, p_month, options.p_week, p_week, options.p_day, p_day, options.p_hour, p_hour,
                options.world_info, options.auxiliar_files_path, grid)
        else:
            self.temporal = None
            settings.write_log('\t\tNone temporal profile set.', level=2)
        self.temporal_factors = None

        # Creating Speciation Object
        if p_speciation is not None:
            self.speciation = Speciation(p_speciation, options.p_speciation, options.molecular_weights)
        else:
            self.speciation = None
            settings.write_log('\t\tNone speciation profile set.', level=2)

        self.vertical_weights = None

        self.emissions = []

        settings.write_time('EmissionInventory', 'Init', timeit.default_timer() - st_time, level=3)

    def create_pollutants_dicts(self, pollutants):
        """
        Create a list of dictionaries with the information of the name, paht and Dataset of each pollutant

        :param pollutants: List of pollutants names
        :type pollutants: list

        :return: List of dictionaries
        :rtype: list
        """

        pollutant_list = []

        for pollutant_name in pollutants:
            pollutant_list.append(
                {'name': pollutant_name,
                 'path': self.get_input_path(pollutant_name),
                 'Dataset': "{0}_{1}".format(self.inventory_name, self.sector)}
            )
        return pollutant_list

    @staticmethod
    def get_profile(id_aux):
        """
        Parse the id of the profiles.

        :param id_aux: ID of the profile.
        :type id_aux: str

        :return: ID of the profile parsed.
        :rtype: str
        """
        import pandas as pd

        if pd.isnull(id_aux):
            return None
        else:
            return id_aux

    def get_input_path(self, pollutant=None):
        """
        Completes the path of the input file that contains the needed information of the given pollutant.

        :param pollutant: Name of the pollutant.
        :type pollutant: str

        :return: Full path of the needed file.
        :rtype: str
        """
        import pandas as pd

        if self.source_type == 'area':
            extension = 'nc'

        elif self.source_type == 'point':
            if self.inventory_name == 'GFASv12':
                extension = 'nc'
            else:
                extension = 'csv'
        else:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: Unknown source type {0}'.format(self.source_type))
            sys.exit(1)

        # Finding upper folder
        if pd.isnull(self.sector):
            upper_folder = '{0}'.format(pollutant)
        else:
            upper_folder = '{0}_{1}'.format(pollutant, self.sector)

        # Finding pollutant folder and filename.
        if self.input_frequency == 'yearly':
            file_name = "{0}_{1}.{2}".format(pollutant, self.reference_year, extension)
        elif self.input_frequency == 'monthly':
            file_name = "{0}_{1}{2}.{3}".format(pollutant, self.reference_year, self.date.strftime("%m"), extension)
        elif self.input_frequency == 'daily':
            file_name = "{0}_{1}{2}{3}.{4}".format(pollutant, self.reference_year, self.date.strftime("%m"),
                                                   self.date.strftime("%d"), extension)
        else:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise ValueError(
                    "ERROR: frequency {0} not implemented. Use yearly, monthly or daily.".format(self.input_frequency))
            sys.exit(1)

        # Filename
        file_path = os.path.join(self.inputs_path, upper_folder, file_name)

        # Checking input file
        if not os.path.exists(file_path):
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise IOError('ERROR: File {0} not found.'.format(file_path))
            sys.exit(1)

        return file_path

    def do_regrid(self):

        st_time = timeit.default_timer()

        settings.write_log("\tRegridding", level=2)
        regridded_emissions = self.regrid.start_regridding()
        for emission in regridded_emissions:
            dict_aux = {'name': emission['name'], 'data': emission['data'], 'units': 'm'}
            self.emissions.append(dict_aux)
        settings.write_time('EmissionInventory', 'do_regrid', timeit.default_timer() - st_time, level=2)

    @staticmethod
    def make_emission_list(options, grid, vertical_output_profile, date):
        """
        Extract the information of the cross table to read all the needed emissions.

        :param options: Full list of parameters given by passing argument or in the configuration file.
        :type options: Namespace

        :param grid: Grid to use.
        :type grid: Grid

        :param vertical_output_profile: Path to eht file that contains the vertical profile.
        :type vertical_output_profile: str

        :param date: Date to simulate.
        :type date: datetime.datetime

        :return: List of Emission inventories already loaded.
        :rtype: list of EmissionInventory
        """
        import pandas as pd
        import re
        from point_gfas_emission_inventory import PointGfasEmissionInventory
        from gfas_emission_inventory import GfasEmissionInventory
        from point_source_emission_inventory import PointSourceEmissionInventory

        st_time = timeit.default_timer()
        settings.write_log('Loading emissions')

        path = options.cross_table
        df = pd.read_csv(path, sep=';', index_col=False)
        for column in ['ei', 'sector', 'ref_year', 'active', 'factor_mask', 'regrid_mask', 'pollutants', 'path',
                       'frequency', 'source_type', 'p_vertical', 'p_month', 'p_week', 'p_day', 'p_hour',
                       'p_speciation']:
            df_cols = list(df.columns.values)
            if column not in df_cols:
                settings.write_log('ERROR: Check the .err file to get more info.')
                if settings.rank == 0:
                    raise AttributeError('ERROR: Column {0} is not in the {1} file.'.format(column, path))
                sys.exit(1)

        df = df[df['active'] == 1]
        num = 1
        emission_inventory_list = []
        for i, emission_inventory in df.iterrows():
            settings.write_log('\tLoading emission {0}/{1} (Inventory: {2}; Sector: {3})'.format(
                num, len(df), emission_inventory.ei, emission_inventory.sector), level=1)
            num += 1
            pollutants = list(map(str, re.split(', |,|; |;| ', emission_inventory.pollutants)))

            # MONTHLY
            try:
                # gridded temporal profile
                p_month = emission_inventory.p_month.replace('<data_path>', options.data_path)
            except AttributeError:
                p_month = emission_inventory.p_month
            # WEEKLY
            try:
                # gridded temporal profile
                p_week = emission_inventory.p_week.replace('<data_path>', options.data_path)
            except AttributeError:
                p_week = emission_inventory.p_week
            # DAILY
            try:
                # gridded temporal profile
                p_day = emission_inventory.p_day.replace('<data_path>', options.data_path)
            except AttributeError:
                p_day = emission_inventory.p_day
            # HOURLY
            try:
                # gridded temporal profile
                p_hour = emission_inventory.p_hour.replace('<data_path>', options.data_path)
            except AttributeError:
                p_hour = emission_inventory.p_hour

            emission_inventory_path = emission_inventory.path.replace('<data_path>', options.data_path)
            emission_inventory_path = emission_inventory_path.replace('<input_dir>', options.input_dir)

            if emission_inventory.source_type == 'area':
                if emission_inventory.ei == 'GFASv12':
                    emission_inventory_list.append(
                        GfasEmissionInventory(options, grid, date, emission_inventory.ei,
                                              emission_inventory.source_type, emission_inventory.sector, pollutants,
                                              emission_inventory_path,
                                              emission_inventory.frequency, vertical_output_profile,
                                              reference_year=emission_inventory.ref_year,
                                              factors=emission_inventory.factor_mask,
                                              regrid_mask=emission_inventory.regrid_mask,
                                              p_vertical=emission_inventory.p_vertical,
                                              p_month=p_month,
                                              p_week=p_week,
                                              p_day=p_day,
                                              p_hour=p_hour,
                                              p_speciation=emission_inventory.p_speciation))
                else:
                    emission_inventory_list.append(
                        EmissionInventory(options, grid, date, emission_inventory.ei, emission_inventory.source_type,
                                          emission_inventory.sector, pollutants,
                                          emission_inventory_path,
                                          emission_inventory.frequency, vertical_output_profile,
                                          reference_year=emission_inventory.ref_year,
                                          factors=emission_inventory.factor_mask,
                                          regrid_mask=emission_inventory.regrid_mask,
                                          p_vertical=emission_inventory.p_vertical,
                                          p_month=p_month,
                                          p_week=p_week,
                                          p_day=p_day,
                                          p_hour=p_hour,
                                          p_speciation=emission_inventory.p_speciation))
            elif emission_inventory.source_type == 'point':
                if emission_inventory.ei == 'GFASv12':
                    emission_inventory_list.append(
                        PointGfasEmissionInventory(
                            options, grid, date, emission_inventory.ei, emission_inventory.source_type,
                            emission_inventory.sector, pollutants, emission_inventory_path,
                            emission_inventory.frequency, vertical_output_profile,
                            reference_year=emission_inventory.ref_year, factors=emission_inventory.factor_mask,
                            regrid_mask=emission_inventory.regrid_mask, p_vertical=emission_inventory.p_vertical,
                            p_month=p_month, p_week=p_week, p_day=p_day, p_hour=p_hour,
                            p_speciation=emission_inventory.p_speciation))
                else:
                    emission_inventory_list.append(
                        PointSourceEmissionInventory(options, grid, date, emission_inventory.ei,
                                                     emission_inventory.source_type, emission_inventory.sector,
                                                     pollutants, emission_inventory_path,
                                                     emission_inventory.frequency, vertical_output_profile,
                                                     reference_year=emission_inventory.ref_year,
                                                     factors=emission_inventory.factor_mask,
                                                     regrid_mask=emission_inventory.regrid_mask,
                                                     p_vertical=emission_inventory.p_vertical,
                                                     p_month=p_month,
                                                     p_week=p_week,
                                                     p_day=p_day,
                                                     p_hour=p_hour,
                                                     p_speciation=emission_inventory.p_speciation))
            else:
                settings.write_log('ERROR: Check the .err file to get more info.')
                if settings.rank == 0:
                    raise ValueError("ERROR: The emission inventory source type '{0}'".format(
                        emission_inventory.source_type) +
                                     " is not implemented. Use 'area' or 'point'")
                sys.exit(1)
            settings.write_log('', level=2)
        settings.write_time('EmissionInventory', 'make_emission_list', timeit.default_timer() - st_time, level=3)

        return emission_inventory_list


if __name__ == "__main__":
    pass
