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

import hermesv3_gr.config.settings as settings
from emission_inventory import EmissionInventory


class GfasEmissionInventory(EmissionInventory):
    """
    Class that defines the content and the methodology for the GFAS 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 frequency: Frequency of the inputs. [yearly, monthly, daily]
    :type frequency: str

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

    :param factors: NOT IMPLEMENTED YET
    :type factors: NOT IMPLEMENTED YET

    :param regrid_mask: NOT IMPLEMENTED YET
    :type regrid_mask: NOT IMPLEMENTED YET

    :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,
                 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.vertical.vertical_gfas import GfasVerticalDistribution

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

        super(GfasEmissionInventory, self).__init__(
            options, grid, current_date, inventory_name, source_type, sector, pollutants, inputs_path, frequency,
            vertical_output_profile,
            reference_year=reference_year, factors=factors, regrid_mask=regrid_mask, p_vertical=None,
            p_month=p_month, p_week=p_week, p_day=p_day, p_hour=p_hour, p_speciation=p_speciation)

        self.approach = self.get_approach(p_vertical)
        self.method = self.get_method(p_vertical)

        self.altitude = self.get_altitude()

        self.vertical = GfasVerticalDistribution(vertical_output_profile, self.approach, self.get_altitude())

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

    def get_input_path(self, pollutant=None, extension='nc'):
        """
        Completes the path of the NetCDF that contains the needed information of the given pollutant.

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

        :param extension: Extension of the input file.
        :type: str

        :return: Full path of the needed NetCDF.
        :rtype: str
        """
        st_time = timeit.default_timer()

        netcdf_path = os.path.join(self.inputs_path, 'multivar', 'ga_{0}.{1}'.format(
            self.date.strftime('%Y%m%d'), extension))

        settings.write_time('GfasEmissionInventory', 'get_input_path', timeit.default_timer() - st_time, level=3)

        return netcdf_path

    def get_altitude(self):
        """
        Extract the altitude values depending on the choosen method.

        :return: Array with the alittude of each fire.
        :rtype: numpy.array
        """
        from hermesv3_gr.tools.netcdf_tools import extract_vars

        st_time = timeit.default_timer()

        if self.method == 'sovief':
            alt_var = 'apt'
        elif self.method == 'prm':
            alt_var = 'mami'
        else:
            alt_var = None

            print "ERROR: Only 'sovief' and 'prm' methods are accepted."

        [alt] = extract_vars(self.get_input_path(), [alt_var])

        alt = alt['data']

        settings.write_time('GfasEmissionInventory', 'get_altitude', timeit.default_timer() - st_time, level=3)
        return alt

    @ staticmethod
    def get_approach(p_vertical):
        """
        Extract the given approach value.

        :return: Approach value
        :rtype: str
        """
        import re

        st_time = timeit.default_timer()

        return_value = None
        aux_list = re.split(', |,| , | ,', p_vertical)
        for element in aux_list:
            aux_value = re.split('=| =|= | = ', element)
            if aux_value[0] == 'approach':
                return_value = aux_value[1]

        settings.write_time('GfasEmissionInventory', 'get_approach', timeit.default_timer() - st_time, level=3)

        return return_value

    @ staticmethod
    def get_method(p_vertical):
        """
        Extract the given method value.

        :return: Method value
        :rtype: str
        """
        import re

        st_time = timeit.default_timer()

        return_value = None
        aux_list = re.split(', |,| , | ,', p_vertical)
        for element in aux_list:
            aux_value = re.split('=| =|= | = ', element)
            if aux_value[0] == 'method':
                return_value = aux_value[1]

        settings.write_time('GfasEmissionInventory', 'get_method', timeit.default_timer() - st_time, level=3)

        return return_value

    def do_vertical_allocation(self, values):
        """
        Allocates the fire emissions on their top level.

        :param values: 2D array with the fire emissions
        :type values: numpy.array

        :return: Emissions already allocated on the top altitude of each fire.
        :rtype: numpy.array
        """
        st_time = timeit.default_timer()

        return_value = self.vertical.do_vertical_interpolation_allocation(values, self.altitude)

        settings.write_time('GfasEmissionInventory', 'do_vertical_allocation', timeit.default_timer() - st_time,
                            level=3)

        return return_value

    def do_regrid(self):

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

        for i in xrange(len(self.emissions)):
            self.emissions[i]["data"] = self.do_vertical_allocation(self.emissions[i]["data"])

        regridded_emissions = self.regrid.start_regridding(gfas=True, vertical=self.vertical)

        for emission in regridded_emissions:
            dict_aux = {'name': emission['name'], 'data': emission['data'], 'units': 'm'}
            # dict_aux['data'] = dict_aux['data'].reshape((1,) + dict_aux['data'].shape)
            self.emissions.append(dict_aux)
        self.vertical = None

        settings.write_time('GfasEmissionInventory', 'do_regrid', timeit.default_timer() - st_time, level=2)


if __name__ == "__main__":
    pass
