#!/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 re
import numpy as np
from nes import open_netcdf
from hermesv3_gr.modules.masking.masking import Masking
from hermesv3_gr.config import precision, log_message
from .emission_inventory import EmissionInventory


class GfasEmissionInventory(EmissionInventory):

    def __init__(self, options, grid, current_date, inventory_name, source_type, sector, pollutants, inputs_path,
                 frequency, vertical_output_profile,
                 reference_year=2010, coverage='global', 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.approach = self.parse_profile_gfas(p_vertical, 'approach')
        self.method = self.parse_profile_gfas(p_vertical, 'method')
        self.altitude = self.get_altitude()
        self.layer_widths = grid.lev['data'].copy()
        self.layer_widths[1:] -= self.layer_widths[:-1]

        super(GfasEmissionInventory, self).__init__(
            options, grid, current_date, inventory_name, source_type, sector, pollutants, inputs_path, frequency,
            vertical_output_profile,
            reference_year=reference_year, coverage=coverage, 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,
            countries_shapefile=countries_shapefile)

        self.vertical = None

    @staticmethod
    def print_ini_message():
        """
        Print initial message
        """
        log_message('Creating GFAS area source emission inventory.', level=3)
        return None
    
    def prepare_input_emissions(self, write_aux_file=True):
        """
        Prepare the emissions but without writing the auxiliary file.

        1. Combine the pollutants of the source
        2. Apply the masks
        2. Perform the Regrid
        3. DO NOT save the regridded file as auxiliary files

        Parameters
        ----------
        write_aux_file : bool
            Indicates if you want to save the regridded emissions as auxiliary file.
        """

        super(GfasEmissionInventory, self).prepare_input_emissions(write_aux_file=False)

        return self.emissions

    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`
        """

        # In the initialization only the first pollutant is read but not loaded.
        self.emissions = open_netcdf(self.src_emis_file_list[0], comm=self.grid.comm, balanced=self.grid.balanced,
                                     parallel_method=self.grid.parallel_method)
        self.emissions.keep_vars(self.input_pollutants + [self.altitude])
        
        # self.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())
        self.emissions.load()

        # Casting emissions to the selected precision
        self.emissions.to_dtype(precision)

        # Masking
        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=self.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:
            for var_name in self.emissions.variables.keys():
                self.emissions.variables[var_name]['data'] *= mask_factors

        # Vertical distribution
        vertical_factor = self.get_vertical_weights(self.emissions.variables[self.altitude]['data'])
        self.emissions.free_vars(self.altitude)

        for var_name in self.emissions.variables.keys():
            self.emissions.variables[var_name]['data'] = np.array(
                self.emissions.variables[var_name]['data'] * vertical_factor,
                dtype=precision)

        return None

    def get_src_file_list(self, pollutant_list):
        """
        Generates the input file list

        Parameters
        ----------
        pollutant_list : list
            Pollutant list

        Returns
        -------
        List[str]
            List of source file paths
        """
        file_list = [self.get_input_path(pollutant=None)]
        return file_list

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

        Parameters
        ----------
        pollutant : str , None
            Name of the pollutant of the NetCDF.
        extension : str
            Extension of the input file.

        Returns
        -------
        str
            Full path of the needed NetCDF.
        """
        netcdf_path = os.path.join(self.inputs_path, 'multivar', 'ga_{0}.{1}'.format(
            self.date.strftime('%Y%m%d'), extension))
        return netcdf_path

    def get_altitude(self):
        """
        Obtain the altitude variable name depending on the selected method

        Returns
        -------
        str
            Altitude variable name
        """

        if self.method == 'sovief':
            alt_var = 'apt'
        elif self.method == 'prm':
            alt_var = 'mami'
        else:
            raise ValueError("ERROR: Only 'sovief' and 'prm' methods are accepted.")

        return alt_var

    @ staticmethod
    def parse_profile_gfas(p_vertical, value):
        """
        Parses the vertical profile to obtain the corresponding value

        Parameters
        ----------
        p_vertical : str
            Profile to parse
        value : str
            Value to extract
            Options: 'approach' or 'method'
        Returns
        -------
        str
            Extracted value
        """
        return_value = None
        aux_list = re.split(', |,| , | ,', p_vertical)
        for element in aux_list:
            aux_value = re.split('=| =|= | = ', element)
            if aux_value[0] == value:
                return_value = aux_value[1]

        return return_value

    def get_vertical_weights(self, altitude_data):
        """
        Compute the input vertical weights

        Parameters
        ----------
        altitude_data : np.ndarray
            2D (4D with 1st and 2nd dim as 1) with the altitude of each emission

        Returns
        -------
        np.array
            3D (4D with first dim as 1) with the vertical distribution.
        """
        weights_3d = np.zeros((1, len(self.grid.lev['data']), altitude_data.shape[-2], altitude_data.shape[-1]),
                              dtype=precision)

        top_layer = self.get_top_layer_data(altitude_data)

        for y in range(altitude_data.shape[-2]):
            for x in range(altitude_data.shape[-1]):
                weights_3d[0, :, y, x] = self.get_cell_weight_list(altitude_data[0, 0, y, x], top_layer[y, x])

        return weights_3d

    def get_top_layer_data(self, altitude_data):
        """
        Obtain a 2D array with the top vertical layer to allocate the emissions

        Parameters
        ----------
        altitude_data : np.ndarray
            Array with the altitude of the emission

        Returns
        -------
        np.ndarray
            2D array with the top vertical layer where allocate the emissions.
        """
        aux_altitude = altitude_data.copy()
        top_data = np.zeros(aux_altitude.shape, dtype=int)

        for level_i, level_m in enumerate(self.grid.lev['data']):
            top_data[aux_altitude <= level_m] = level_i
            aux_altitude[aux_altitude <= level_m] = None
        del aux_altitude
        top_data = top_data.reshape((altitude_data.shape[-2], altitude_data.shape[-1]))

        return top_data

    def get_cell_weight_list(self, altitude, top_layer):
        """
        Obtain the vertical distribution of a single point.

        Parameters
        ----------
        altitude : float
            Altitude variable.
        top_layer : int
            Top layer of the emission.

        Returns
        -------
        np.array
            1D array (with vertical level length) with the weight of emission to distribute vertically of a single
            cell emission.
        """
        vertical_weights = np.zeros((len(self.grid.lev['data'])), dtype=precision)
        if altitude <= self.grid.lev['data'][0]:
            vertical_weights[0] = 1
        else:
            if self.approach == '50_top':
                vertical_weights[top_layer] = 0.5
                to_distribute = 0.5
                top_layer -= 1
                # Changing altitude to the new top vertical layer
                altitude = self.grid.lev['data'][top_layer]

            elif self.approach == 'uniform':
                to_distribute = 1.
            else:
                raise ValueError("Unknown {0} GFAS vertical approach".format(self.approach))

            for level_i in range(top_layer + 1):
                vertical_weights[level_i] = to_distribute * self.layer_widths[level_i] / altitude

        return vertical_weights
