#!/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 hermesv3_gr.config.settings as settings
from .emission_inventory import EmissionInventory


class PointSourceEmissionInventory(EmissionInventory):
    """
    Class that defines the content and the methodology for the Point Source emission inventories

    :param options: Place where are stored all the arguments.
    :type options: namespace

    :param grid: Destination grid object.
    :type grid: Grid

    :param current_date: Date of the day 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: Description of the scale factors per country. (e.g. Spain 1.5, China 3.)
    :type factors: str

    :param regrid_mask: Description of the masking countries (adding e.g. + Spain Andorra) (subtracting e.g. - Spain)
    :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, comm, 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,
                 countries_shapefile=None):

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

        super(PointSourceEmissionInventory, self).__init__(
            comm, 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=p_vertical, p_month=p_month, p_week=p_week, p_day=p_day, p_hour=p_hour,
            p_speciation=p_speciation, countries_shapefile=countries_shapefile)

        self.crs = {'init': 'epsg:4326'}
        self.location = None
        self.area = None
        self.vertical = 'custom'

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

    def do_regrid(self):
        """
        Allocates the point source emission on the correspondent cell (getting the ID of the cell).

        :return: True when everything is correct.
        :rtype: bool
        """
        import pandas as pd
        import geopandas as gpd
        from shapely.geometry import Point

        st_time = timeit.default_timer()
        settings.write_log("\tAllocating point sources on grid:", level=2)

        num = 1
        for pollutant in self.pollutant_dicts:
            if self.location is None:
                grid_shape = self.grid.to_shapefile(full_grid=False)

            settings.write_log('\t\tPollutant {0} ({1}/{2})'.format(
                pollutant['name'], num, len(self.pollutant_dicts)), level=3)
            num += 1

            df = pd.read_csv(pollutant['path'])

            geometry = [Point(xy) for xy in zip(df.Lon, df.Lat)]
            df = gpd.GeoDataFrame(df.loc[:, ['Emis', 'Alt_Injection']], crs=self.crs, geometry=geometry)

            df = df.to_crs(grid_shape.crs)
            df = gpd.sjoin(df, grid_shape, how="inner", op='intersects')

            # Drops duplicates when the point source is on the boundary of the cell
            df = df[~df.index.duplicated(keep='first')]

            if self.location is None:
                self.location = df.loc[:, ['Alt_Injection', 'FID']]
                self.area = self.grid.cell_area.flatten()[self.location['FID'].values]

            dict_aux = {
                'name': pollutant['name'],
                'units': '...',
                'data': df.loc[:, 'Emis'].values / self.area
            }

            self.emissions.append(dict_aux)
        settings.write_time('PointSourceEmissionInventory', 'do_regrid', timeit.default_timer() - st_time, level=2)
        return True

    def calculate_altitudes(self, vertical_description_path):
        """
        Calculate the number layer to allocate the point source.

        :param vertical_description_path: Path to the file that contains the vertical description
        :type vertical_description_path: str

        :return: True
        :rtype: bool
        """
        import pandas as pd

        st_time = timeit.default_timer()
        settings.write_log("\t\tCalculating vertical allocation.", level=3)
        df = pd.read_csv(vertical_description_path, sep=',')
        # df.sort_values(by='height_magl', ascending=False, inplace=True)
        self.location['layer'] = None

        for i, line in df.iterrows():
            self.location.loc[self.location['Alt_Injection'] <= line['height_magl'], 'layer'] = line['Ilayer'] - 1
            self.location.loc[self.location['Alt_Injection'] <= line['height_magl'], 'Alt_Injection'] = None
        del self.location['Alt_Injection']

        settings.write_time('PointSourceEmissionInventory', 'calculate_altitudes', timeit.default_timer() - st_time,
                            level=2)

        return True

    def point_source_by_cell(self):
        """
        Sums the different emissions that are allocated in the same cell and layer.

        :return: None
        """
        st_time = timeit.default_timer()

        aux_df = None
        for emission in self.emissions:

            aux_df = self.location.copy()
            aux_df['Emis'] = emission['data']
            aux_df = aux_df.groupby(['FID', 'layer']).sum()
            aux_df.reset_index(inplace=True)
            emission['data'] = aux_df['Emis']

        self.location = aux_df.loc[:, ['FID', 'layer']]

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

        return None


if __name__ == "__main__":
    pass
