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

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


class PointGfasEmissionInventory(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 as point source.', level=3)
        super(PointGfasEmissionInventory, 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.get_approach(p_vertical),
                                                 self.get_altitude())

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

    def __str__(self):
        string = "PointGfasEmissionInventory:\n"
        string += "\t self.method = {0}\n".format(self.method)
        string += "\t self.vertical = {0}\n".format(self.vertical)

        return string

    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()

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

        # 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)

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

        return file_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('PointGfasEmissionInventory', '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('PointGfasEmissionInventory', '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('PointGfasEmissionInventory', '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
        """
        print 'do_vertical_allocation'
        sys.exit()
        st_time = timeit.default_timer()

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

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

        return return_value

    def do_regrid(self):
        import pandas as pd
        import geopandas as gpd
        from shapely.geometry import Point
        import numpy as np
        from netCDF4 import Dataset

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

        netcdf = Dataset(self.pollutant_dicts[0]['path'], mode='r')
        gdf = gpd.GeoDataFrame(crs={'init': 'epsg:4326'})

        # gdf['src_index'] = np.where(self.vertical.altitude.flatten() > 0)[0]
        # print self.input_pollutants[0]
        first_var = netcdf.variables[self.input_pollutants[0]][:]
        # first_var = netcdf.variables['bc'][:]

        # print first_var
        gdf['src_index'] = np.where(first_var.flatten() > 0)[0]

        # Masking
        if self.masking.regrid_mask is not None:
            gdf = gdf.loc[gdf['src_index'].isin(np.where(self.masking.regrid_mask.flatten() > 0)[0]), ]

        gdf['altitude'] = self.vertical.altitude.flatten()[gdf['src_index']]

        lat_1d = netcdf.variables['lat'][:]
        lon_1d = netcdf.variables['lon'][:]
        lon_1d[lon_1d > 180] -= 360
        # 1D to 2D
        lats = np.array([lat_1d] * len(lon_1d)).T
        lons = np.array([lon_1d] * len(lat_1d))

        gdf['geometry'] = [Point(xy) for xy in zip(lons.flatten()[gdf['src_index']],
                                                   lats.flatten()[gdf['src_index']])]
        gdf['src_area'] = netcdf.variables['cell_area'][:].flatten()[gdf['src_index']]
        grid_shp = self.grid.to_shapefile(full_grid=False)

        temp_coords = Dataset(os.path.join(self.grid.temporal_path, 'temporal_coords.nc'), mode='r')

        cell_area = temp_coords.variables['cell_area'][self.grid.x_lower_bound:self.grid.x_upper_bound,
                                                       self.grid.y_lower_bound:self.grid.y_upper_bound].flatten()
        grid_shp['dst_area'] = cell_area[grid_shp['FID']]

        gdf = gpd.sjoin(gdf.to_crs(grid_shp.crs), grid_shp, how='inner')

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

            aux = netcdf.variables[pollutant['name']][:].flatten()[gdf['src_index']]

            if self.masking.scale_mask is not None:
                aux = aux * self.masking.scale_mask.flatten()[gdf['src_index']]

            gdf[pollutant['name']] = aux * gdf['src_area']
            # print 'masa {0}: {1} '.format(pollutant['name'], gdf[pollutant['name']].sum())
            gdf[pollutant['name']] = (aux / gdf['dst_area']) * \
                netcdf.variables['cell_area'][:].flatten()[gdf['src_index']]
        # print netcdf.variables['bc'][:].sum()

        netcdf.close()

        settings.write_time('PointGfasEmissionInventory', 'do_regrid', timeit.default_timer() - st_time, level=2)
        del gdf['src_index'], gdf['index_right']
        # print gdf.sum().to_csv('/home/Earth/ctena/temp/gdf_serie.csv'.format(settings.rank))
        # sys.exit(1)
        self.emissions = gdf

        return True

    def do_regrid_balanced_NOT_USED(self):
        import pandas as pd
        import geopandas as gpd
        from shapely.geometry import Point
        import numpy as np
        from netCDF4 import Dataset

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

        netcdf = Dataset(self.pollutant_dicts[0]['path'], mode='r')
        if settings.rank == 0:
            gdf = gpd.GeoDataFrame(crs={'init': 'epsg:4326'})

            first_var = netcdf.variables[self.input_pollutants[0]][:]

            gdf['src_index'] = np.where(first_var.flatten() > 0)[0]

            gdf['altitude'] = self.vertical.altitude.flatten()[gdf['src_index']]

            lat_1d = netcdf.variables['lat'][:]
            lon_1d = netcdf.variables['lon'][:]
            lon_1d[lon_1d > 180] -= 360
            # 1D to 2D
            lats = np.array([lat_1d] * len(lon_1d)).T
            lons = np.array([lon_1d] * len(lat_1d))

            gdf['geometry'] = [Point(xy) for xy in zip(lons.flatten()[gdf['src_index']],
                                                       lats.flatten()[gdf['src_index']])]
            gdf['src_area'] = netcdf.variables['cell_area'][:].flatten()[gdf['src_index']]
            grid_shp = self.grid.to_shapefile()

            temp_coords = Dataset(os.path.join(self.grid.temporal_path, 'temporal_coords.nc'), mode='r')
            cell_area = temp_coords.variables['cell_area'][:].flatten()
            grid_shp['dst_area'] = cell_area[grid_shp['FID']]
            # grid_shp['dst_area'] = grid_shp.to_crs({'init': 'epsg:3857'}).area
            # print grid_shp.crs['units']
            # sys.exit()

            gdf = gpd.sjoin(gdf.to_crs(grid_shp.crs), grid_shp, how='inner')
            print gdf
            gdf = np.array_split(gdf, settings.size)
        else:
            gdf = None

        gdf = settings.comm.scatter(gdf, root=0)

        for num, pollutant in enumerate(self.pollutant_dicts):
            settings.write_log('\t\tPollutant {0} ({1}/{2})'.format(
                pollutant['name'], num + 1, len(self.pollutant_dicts)), level=3)
            print ('\t\tPollutant {0} ({1}/{2})'.format(pollutant['name'], num + 1, len(self.pollutant_dicts)))
            aux = netcdf.variables[pollutant['name']][:].flatten()[gdf['src_index']]
            gdf[pollutant['name']] = aux * gdf['src_area']
            # print 'masa {0}: {1} '.format(pollutant['name'], gdf[pollutant['name']].sum())
            gdf[pollutant['name']] = (aux / gdf['dst_area']) * netcdf.variables['cell_area'][:].flatten()[
                gdf['src_index']]
        # print netcdf.variables['bc'][:].sum()

        netcdf.close()

        settings.write_time('PointGfasEmissionInventory', 'do_regrid', timeit.default_timer() - st_time, level=2)
        print 'regrid done'
        del gdf['src_index'], gdf['index_right']
        self.emissions = gdf

        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=';')
        if self.vertical.approach == 'surface':
            self.emissions['altitude'] = 0

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

        # Fires with higher altitudes than the max layer limit goes to the last layer
        if len(self.emissions[~self.emissions['altitude'].isna()]) > 0:
            self.emissions.loc[~self.emissions['altitude'].isna(), 'layer'] = layer
            warnings.warn('WARNING: One or more fires have an altitude of fire emission injection higher than the top' +
                          ' layer of the model defined in the {0} file'.format(vertical_description_path))

            print self.emissions.loc[~self.emissions['altitude'].isna()]
        del self.emissions['altitude']

        self.emissions = self.emissions.groupby(['FID', 'layer']).sum()
        self.emissions.reset_index(inplace=True)
        # print '---> Rank {0}, {1}, {2}'.format(settings.rank, self.emissions, len(self.emissions))
        # sys.exit()
        self.emissions = self.vertical.distribute_vertically(self.emissions, self.input_pollutants)

        settings.write_time('PointGfasEmissionInventory', '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
        """
        emissions = self.emissions
        self.location = emissions.loc[:, ['FID', 'layer']]

        self.emissions = []
        for input_pollutant in self.input_pollutants:
            dict_aux = {
                'name': input_pollutant,
                'units': '...',
                'data': emissions.loc[:, input_pollutant].values
            }
            self.emissions.append(dict_aux)


if __name__ == "__main__":
    pass
