#!/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
from datetime import datetime
import pandas as pd
import hermesv3_gr.config.settings as settings
from .emission_inventory import EmissionInventory
from hermesv3_gr.modules.temporal.temporal import TemporalDistribution
from hermesv3_gr.modules.vertical.vertical_gfas_hourly import GfasHourlyVerticalDistribution
from hermesv3_gr.tools.mpi_tools import *
# import mpi4py
# mpi4py.rc.recv_mprobe = False

BUFFER_SIZE = 2**25


class PointGfasHourlyEmissionInventory(EmissionInventory):
    """
    Class that defines the content and the methodology for the GFAS hourly 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_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_speciation=None, countries_shapefile=None):

        st_time = timeit.default_timer()
        settings.write_log('\t\tCreating GFAS emission inventory as point source.', level=3)
        super(PointGfasHourlyEmissionInventory, 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=None, p_month=None,
            p_week=None, p_day=None, p_hour=None, p_speciation=p_speciation, countries_shapefile=countries_shapefile)

        self.method = self.get_method(p_vertical)
        self.altitude_name = self.get_altitude_name()

        self.vertical = GfasHourlyVerticalDistribution(vertical_output_profile, self.get_approach(p_vertical))

        self.date_array = TemporalDistribution.caclulate_date_array(
            options.start_date, options.output_timestep_type, options.output_timestep_num, options.output_timestep_freq)

        self.grid_shp = self.grid.to_shapefile(full_grid=False)
        self.cell_id_by_proc = self.get_cell_ids()
        self.coordinates = self.get_coordinates()
        self.fid_distribution = self.get_fids()

        self.temporal = None

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

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

        return string

    @ 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('PointGfasHourlyEmissionInventory', 'get_approach', timeit.default_timer() - st_time,
                            level=3)

        return return_value

    def get_coordinates(self):
        from netCDF4 import Dataset
        import geopandas as gpd
        import pandas as pd
        from shapely import wkt
        from hermesv3_gr.tools.netcdf_tools import get_grid_area

        # Getting 2D coordinates
        coordinates_path = os.path.join(
            self.auxiliary_files_path, 'gfas', 'coordinates_s{0}_n{1}.csv'.format(settings.size, settings.rank))

        if not os.path.exists(coordinates_path):
            if settings.rank == 0:
                if not os.path.exists(os.path.dirname(coordinates_path)):
                    os.makedirs(os.path.dirname(coordinates_path), exist_ok=True)

                coordinates_full = os.path.join(
                    self.auxiliary_files_path, 'gfas', 'coordinates_full.csv'.format(settings.size, settings.rank))
                if not os.path.exists(coordinates_full):
                    settings.write_log('\t\tCreating GFAS coordinates shapefile. It may take a few minutes.', level=3)
                    src_area = get_grid_area(self.get_input_path(date=self.date_array[0]))
                    netcdf = Dataset(self.get_input_path(date=self.date_array[0]), mode='r')
                    try:
                        lats = netcdf.variables['latitude'][:]
                    except KeyError:
                        try:
                            lats = netcdf.variables['lat'][:]
                        except KeyError:
                            raise KeyError("Neither 'lat' or 'latitude' are present on the {0} file".format(
                                self.get_input_path(date=self.date_array[0])))
                    try:
                        lons = netcdf.variables['longitude'][:]
                    except KeyError:
                        try:
                            lons = netcdf.variables['lon'][:]
                        except KeyError:
                            raise KeyError("Neither 'lon' or 'longitude' are present on the {0} file".format(
                                self.get_input_path(date=self.date_array[0])))
                    netcdf.close()

                    lons[lons > 180] = lons[lons > 180] - 360.0

                    len_lats = len(lats)
                    len_lons = len(lons)

                    lats = np.array([lats] * len_lons).T.flatten()
                    lons = np.array([lons] * len_lats).flatten()

                    coordinates = gpd.GeoDataFrame(lats, columns=['lats'], crs={'init': 'epsg:4326'})
                    coordinates['lons'] = lons
                    coordinates.index.name = 'FID'

                    coordinates['geometry'] = \
                        'POINT (' + coordinates['lons'].astype(str) + ' ' + coordinates['lats'].astype(str) + ')'
                    coordinates.drop(columns=['lats', 'lons'], inplace=True)
                    coordinates['src_area'] = src_area.flatten()
                    coordinates['geometry'] = coordinates['geometry'].apply(wkt.loads)
                    coordinates.reset_index(inplace=True)
                    # Too time consumer
                    #coordinates.to_file(coordinates_shapefile)
                    settings.write_log('\t\t\tDone!', level=3)

                    grid = self.grid.to_shapefile(full_grid=True).to_crs(coordinates.crs)

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

                    grid['dst_area'] = temp_coords.variables['cell_area'][:].flatten()
                    temp_coords.close()

                    settings.write_log('\t\tAllocating GFAS coordinates into the Grid cells. It may take a few minutes.',
                                       level=3)

                    coordinates = gpd.sjoin(coordinates, grid, how='inner', op='within')
                    settings.write_log('\t\t\tDone!', level=3)
                    coordinates.rename(columns={'FID_left': 'FID'}, inplace=True)

                    coordinates = coordinates[['FID', 'src_area', 'Cell_ID', 'dst_area']]
                    coordinates.to_csv(coordinates_full)
                else:
                    coordinates = pd.read_csv(coordinates_full)

                # for proc in range(settings.size):
                #     if proc > 0:
                #         settings.comm.send(coordinates[coordinates['Cell_ID'].isin(self.cell_id_by_proc[proc])],
                #                            dest=proc)
                # coordinates = coordinates[coordinates['Cell_ID'].isin(self.cell_id_by_proc[0])]

            else:
                settings.write_log('\t\tWaiting to GFAS coordinates shapefile (rank 0). It may take a few minutes.',
                                   level=3)
                # coordinates = settings.comm.recv(source=0)
                coordinates = None
            coordinates = settings.comm.bcast(coordinates, root=0)
            settings.write_log('\t\tGFAS coordinates received.', level=3)

            coordinates = coordinates[coordinates['Cell_ID'].isin(self.cell_id_by_proc[settings.rank])]
            # # coordinates = coordinates[['FID', 'Cell_ID']]
            # temp_coords = Dataset(os.path.join(self.grid.temporal_path, 'temporal_coords.nc'), mode='r')
            #
            # cell_area = pd.DataFrame(temp_coords.variables['cell_area'][:].flatten(), columns=['dst_area'])
            # temp_coords.close()
            # cell_area['Cell_ID'] = cell_area.index
            #
            # coordinates = pd.merge(coordinates, cell_area, on='Cell_ID')
            sys.stdout.flush()
            coordinates.to_csv(coordinates_path)
        else:
            coordinates = pd.read_csv(coordinates_path)
            coordinates = coordinates[['FID', 'src_area', 'Cell_ID', 'dst_area']]
        coordinates.set_index('FID', inplace=True)

        return coordinates

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

        :param date: Date of the NetCDF.
        :type date: datetime

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

        if date is None:
            date = self.date

        # TODO to change path
        file_path = os.path.join(self.inputs_path, 'multivar', 'ftp',
                                 '{0}'.format(date.strftime('%Y%m%d')),
                                 'ga_{0}.{1}'.format(date.strftime('%Y%m%d_%H'), 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('PointGfasHourlyEmissionInventory', 'get_input_path', timeit.default_timer() - st_time,
                            level=3)

        return file_path

    def get_altitude_name(self):
        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.")

        settings.write_time('PointGfasHourlyEmissionInventory', 'get_altitude_name', timeit.default_timer() - st_time,
                            level=3)
        return alt_var

    @ 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('PointGfasHourlyEmissionInventory', 'get_method', timeit.default_timer() - st_time, level=3)

        return return_value

    def get_cell_ids(self):
        cell_id_by_proc = {}
        for proc in range(settings.size):
            if proc == settings.rank:
                aux_fid = self.grid_shp['Cell_ID'].values
            else:
                aux_fid = None
            cell_id_by_proc[proc] = bcast_array(settings.comm, aux_fid, master=proc)

        return cell_id_by_proc

    def get_fids(self):
        fid_by_proc = {}
        for proc in range(settings.size):
            if proc == settings.rank:
                aux_fid = list(np.unique(self.coordinates.index.get_level_values('FID').values))
            else:
                aux_fid = None
            fid_by_proc[proc] = bcast_array(settings.comm, aux_fid, master=proc)

        return fid_by_proc

    def do_regrid(self):

        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)

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

        distribution = None
        # st_time = timeit.default_timer()
        shapefile_list = []

        # One NetCDF for time step (hour)
        for tstep, date in enumerate(self.date_array):
            shapefile_list_hour = []
            netcdf = Dataset(self.get_input_path(date=date), mode='r')
            for pollutant in self.input_pollutants:
                if distribution is None:
                    distribution = get_balanced_distribution(
                        settings.size, netcdf.variables[pollutant].shape)[settings.rank]
                var = netcdf.variables[pollutant][0, distribution['y_min']: distribution['y_max'], :].flatten()

                var_index = np.where(var > 0)[0]
                shapefile_list_hour.append(
                    pd.DataFrame(var[var_index], columns=[pollutant],
                                 index=pd.MultiIndex.from_product([(var_index + distribution['fid_min']).astype(int),
                                                                   [tstep]], names=['FID', 'tstep'])))

            shapefile = pd.concat(shapefile_list_hour, axis=1, sort=False)

            # shapefile = shapefile[sorted(sum(fid_distribution.values(), []))]

            altitude_var = netcdf.variables[self.altitude_name][:].flatten()
            netcdf.close()

            shapefile['altitude'] = altitude_var[shapefile.index.get_level_values('FID')]

            shapefile_list.append(shapefile)
        shapefile = pd.concat(shapefile_list, sort=False)

        # Filtering shapefile with the involved cells
        shapefile = shapefile[shapefile.index.get_level_values('FID').isin(
            sorted(sum(self.fid_distribution.values(), [])))]

        shapefile = balance_dataframe(settings.comm, shapefile)
        shapefile.fillna(0.0, inplace=True)

        self.emissions = shapefile

        return True

    def distribute(self, emissions):
        settings.write_log("\t\tCalculating distribution.", level=3)
        total_emis = settings.comm.gather(emissions, root=0)
        if settings.rank == 0:
            total_emis = pd.concat(total_emis)
        total_emis = settings.comm.bcast(total_emis, root=0)
        emissions = total_emis[total_emis['FID'].isin(sorted(np.unique(self.fid_distribution[settings.rank])))]

        emissions = pd.merge(emissions, self.coordinates.reset_index(), on='FID')

        for pollutant in self.input_pollutants:
            emissions[pollutant] = emissions[pollutant] * (emissions['src_area'] / emissions['dst_area'])
        emissions.drop(columns=['FID', 'src_area', 'dst_area'], inplace=True)

        # Replacing Cell_ID by FID to be able to write
        emissions = pd.merge(emissions, self.grid_shp[['FID', 'Cell_ID']], on='Cell_ID')
        emissions.drop(columns=['Cell_ID'], inplace=True)

        emissions.rename(columns={'Cell_ID': 'FID'}, inplace=True)
        emissions = emissions.groupby(['FID', 'tstep', 'layer']).sum()

        return emissions

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

        del self.emissions['altitude']

        self.emissions = self.emissions.groupby(['FID', 'tstep', 'layer']).sum()
        self.emissions.reset_index(inplace=True)

        self.emissions = self.vertical.distribute_vertically(self.emissions, self.input_pollutants)
        settings.write_log("\t\t\tVertical allocation done!", level=3)
        settings.write_time('PointGfasHourlyEmissionInventory', '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.distribute(self.emissions).reset_index()
        self.location = emissions.loc[:, ['FID', 'tstep', '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
