#!/usr/bin/env python
import sys
# 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 os
from copy import deepcopy
import pandas as pd
from nes import open_netcdf
from hermesv3_gr.modules.masking.masking import Masking
from pandas import DataFrame
import numpy as np
from hermesv3_gr.config import settings, log_message
from .gfas_emission_inventory import GfasEmissionInventory


class PointGfasEmissionInventory(GfasEmissionInventory):
    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, countries_shapefile=None):

        self.dst_fid = None
        self.dst_level = None
        print(f"Rank {grid.rank}: PRE {grid.get_fids().shape}")
        sys.stdout.flush()
        expanded_grid = grid.copy(copy_vars=False)
        expanded_grid.set_communicator(grid.comm)

        for var_name in expanded_grid.variables.keys():
            print(f"{var_name}: {expanded_grid.variables[var_name]['data'].shape}")
        expanded_grid.expand(n_cells=1)

        super(PointGfasEmissionInventory, self).__init__(
            options, expanded_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)

    def get_weight_matrix_path(self):
        return os.path.join(self.auxiliary_files_path, "Weight_Matrix_{0}_point.nc".format(self.inventory_name))

    @staticmethod
    def print_ini_message():
        log_message('Creating GFAS point source emission inventory.', level=3)

        return None

    def read_input_emissions(self):
        """
        Prepare tje emissions for the regridding

        1. Read the emissions
        2. Apply mask
        3. Distribute vertically

        The resultant emissions are stored in `self.emissions`
        """
        log_message("Loading emissions")
        # In the initialization only the first pollutant is read but not loaded.
        emissions = open_netcdf(self.src_emis_file_list[0], comm=self.grid.comm, balanced=self.grid.balanced,
                                parallel_method=self.grid.parallel_method)
        emissions.sel(lat_min=self.grid.get_full_latitudes_boundaries()['data'].min(),
                      lat_max=self.grid.get_full_latitudes_boundaries()['data'].max(),
                      lon_min=self.grid.get_full_longitudes_boundaries()['data'].min(),
                      lon_max=self.grid.get_full_longitudes_boundaries()['data'].max())

        emissions.keep_vars(self.input_pollutants + [self.altitude])
        emissions.load()

        # Flux to mass
        emissions.calculate_grid_area(overwrite=True)
        for var_name in emissions.variables.keys():
            if var_name != self.get_altitude():
                emissions.variables[var_name]['data'] *= emissions.cell_measures["cell_area"]['data']

        emissions.to_dtype(data_type=settings.precision)

        # Vertical distribution
        vertical_factor = self.get_vertical_weights(emissions.variables[self.altitude]['data'])
        vertical_factor = vertical_factor.reshape(
            (vertical_factor.shape[1], vertical_factor.shape[2] * vertical_factor.shape[3]))
        emissions.free_vars(self.altitude)

        # Obtaining the location of the emissions
        # We assume that the first pollutant is representative of the location of the GFAS data.
        # If the first pollutant has data the others has data too. The same in the other side.
        representative_var = self.input_pollutants[0]
        src_index = np.where(emissions.variables[representative_var]['data'].flatten() > 0)[0]

        src_fids = emissions.get_fids().flatten()
        if len(src_index) > 0:
            src_fids = src_fids[src_index]
            result = self.weight_matrix.loc[src_fids].copy()
            for var_name in self.input_pollutants:
                log_message(f"Var name {var_name}")
                log_message(f"{emissions.variables[var_name]['data']}")
                result[var_name] = emissions.variables[var_name]['data'].flatten()[src_index]
                log_message(f"Var name {var_name} loaded")


            result_by_level = []
            for i_level in range(vertical_factor.shape[0]):
                aux_result = result.reset_index().copy()
                aux_result['level'] = i_level
                aux_result['factor'] += vertical_factor[i_level, src_index]
                aux_result = aux_result.loc[aux_result['factor'] > 0]
                aux_result[self.input_pollutants] = aux_result[self.input_pollutants].multiply(
                    aux_result['factor'], axis="index")
                del aux_result['factor'], aux_result['src_fid']

                result_by_level.append(aux_result)
            result = pd.concat(result_by_level)
            result = result.groupby(['dst_fid', 'level']).sum()

        else:
            result = DataFrame(index=pd.MultiIndex(levels=[[], []], codes=[[], []], names=['dst_fid', 'level']),
                               columns=self.input_pollutants)

        self.emissions = result

        return None

    def init_regrid(self, weight_matrix_path):
        """
        Initialise the Weight matrix

        Parameters
        ----------
        weight_matrix_path : str
            Weight matrix auxiliary file path

        Returns
        -------
        DataFrame
            Weight matrix with the src_fid as index and the dst_fid as column
        """
        # To consider save the Dataframe
        st_time = timeit.default_timer()
        emissions = open_netcdf(comm=self.grid.comm, path=self.src_emis_file_list[0], balanced=self.grid.balanced,
                                parallel_method=self.grid.parallel_method)
        emissions.sel(lat_min=self.grid.get_full_latitudes_boundaries()['data'].min(),
                      lat_max=self.grid.get_full_latitudes_boundaries()['data'].max(),
                      lon_min=self.grid.get_full_longitudes_boundaries()['data'].min(),
                      lon_max=self.grid.get_full_longitudes_boundaries()['data'].max())

        if not os.path.exists(weight_matrix_path):
            # Creating Inverse weight matrix to obtain the nearest dst_cell to each grid_cell
            self.grid.interpolate_horizontal(emissions, weight_matrix_path=weight_matrix_path, info=False,
                                             kind='NearestNeighbour', n_neighbours=1, only_create_wm=True)
        # Reopening weight matrix to be read as normal Nes object
        self.grid.comm.Barrier()

        weight_matrix = open_netcdf(comm=self.grid.comm, path=weight_matrix_path, balanced=self.grid.balanced,
                                    parallel_method=self.grid.parallel_method)
        weight_matrix.keep_vars('idx')
        weight_matrix.load()

        weight_matrix = DataFrame(index=weight_matrix.get_fids().flatten(),
                                  data=weight_matrix.variables['idx']['data'].flatten(), columns=['dst_fid'])
        weight_matrix.index.name = 'src_fid'

        # Masking
        weight_matrix['factor'] = 0
        weight_matrix['factor'] = np.array(weight_matrix['factor'], dtype=settings.precision)
        mask_path = os.path.join(os.path.dirname(self.auxiliary_files_path),
                                 '{0}_WorldMask.nc'.format(self.inventory_name))
        mask_factors = Masking(
            dst_grid=self.grid, src_grid_path=emissions, factors_mask_values=self.factors_mask,
            regrid_mask_values=self.regrid_mask, world_mask_file=mask_path,
            countries_shapefile=self.countries_shapefile).mask_factors

        if mask_factors is not None:
            weight_matrix['factor'] = mask_factors.flatten()

        settings.write_time('PointGfasEmissionInventory', 'init_regrid', timeit.default_timer() - st_time, level=3)
        return weight_matrix

    def do_regrid(self):
        """
        Regridding of the emissions
        """

        st_time = timeit.default_timer()

        log_message("Regridding", level=2)

        # Gathering into 0 and broadcasting emissions
        if self.grid.comm.Get_size() == 1:
            emissions = self.emissions
        else:
            emissions = self.grid.comm.gather(self.emissions, root=0)
            if self.grid.comm.Get_rank() == 0:
                emissions = pd.concat(emissions)
                emissions = emissions.groupby(level=('dst_fid', 'level')).sum()
            else:
                emissions = None
            emissions = self.grid.comm.bcast(emissions, root=0)

        # Selecting only involved emissions
        fid_list = self.grid.get_fids().flatten()
        emissions = emissions[np.in1d(emissions.index.get_level_values('dst_fid'), fid_list)]

        # To Nes mode
        self.dst_fid = emissions.index.get_level_values('dst_fid').to_numpy()
        self.dst_level = emissions.index.get_level_values('level').to_numpy()

        self.emissions = self.grid.copy(copy_vars=False)
        self.emissions.variables = {}
        for poll_name in self.input_pollutants:
            self.emissions.variables[poll_name] = {'data': emissions[poll_name].to_numpy(dtype=settings.precision)}

        # Mass to flux
        fids = self.grid.get_fids()
        self.emissions.calculate_grid_area(overwrite=True)
        for var_name in self.emissions.variables.keys():
            if var_name != self.get_altitude():
                for i_data, cell_emission in enumerate(self.emissions.variables[var_name]['data']):
                    y_pos, x_pos = np.where(fids == self.dst_fid[i_data])
                    self.emissions.variables[var_name]['data'][i_data] /= self.emissions.cell_measures["cell_area"]['data'][y_pos, x_pos]
                # self.emissions.variables[var_name]['data'] /= self.emissions.cell_measures["cell_area"]['data']

        settings.write_time('PointGfasEmissionInventory', 'do_regrid', timeit.default_timer() - st_time, level=2)
        return None

    def get_4d_emissions(self):

        emis = self.emissions.copy(copy_vars=True)

        # Point emissions to 3D grid format
        fids = self.grid.get_fids()
        for var_name in self.emissions.variables.keys():
            log_message(f"{var_name}: \n{self.emissions.variables[var_name]['data']}", level=5)
            if (self.emissions.variables[var_name]['data'] is None) or \
                    (isinstance(self.emissions.variables[var_name]['data'], int)):

                emis.variables[var_name]['data'] = 0
            else:
                # Point to 3D
                data = np.zeros((1, len(self.grid.lev['data']), self.grid.lat['data'].shape[0],
                                 self.grid.lon['data'].shape[-1]), dtype=settings.precision)
                for i_data, cell_emission in enumerate(self.emissions.variables[var_name]['data']):
                    y_pos, x_pos = np.where(fids == self.dst_fid[i_data])
                    data[0, self.dst_level[i_data], y_pos[0], x_pos[0]] = cell_emission

                emis.variables[var_name]['data'] = data

                # Temporal disaggregation
                if self.temporal_factors is not None:
                    emis.variables[var_name]['data'] = np.array(
                        emis.variables[var_name]['data'] * self.temporal_factors[:, np.newaxis, :, :],
                        dtype=settings.precision)
        emis.contract(n_cells=1)
        return emis.variables
