#!/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 vertical import VerticalDistribution


class GfasVerticalDistribution(VerticalDistribution):
    """
    Class that contains all the needed information to vertically distribute the fire emissions.

    :param vertical_output_profile: Path to the file that contains the vertical description of the desired output.
    :type vertical_output_profile: str

    :param approach: Approach to take into account.
    :type approach: str
    """
    def __init__(self, vertical_output_profile, approach, altitude):
        st_time = timeit.default_timer()

        self.altitude = altitude
        self.output_heights = vertical_output_profile
        self.approach = approach

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

    def __str__(self):

        string = "GFAS Vertical distribution:\n"

        string += "\t self.approach = {0}\n".format(self.approach)

        string += "\t self.output_heights = {0}\n".format(self.output_heights)

        string += "\t self.altitude = {0}\n".format(self.altitude)

        return string


    @staticmethod
    def calculate_widths(heights_list):
        """
        Calculate the width of each vertical level.

        :param heights_list: List of the top altitude in meters of each level.
        :type heights_list: list

        :return: List with the width of each vertical level.
        :rtype: list
        """
        st_time = timeit.default_timer()

        widths = []
        for i in xrange(len(heights_list)):
            if i == 0:
                widths.append(heights_list[i])
            else:
                widths.append(heights_list[i] - heights_list[i - 1])

        settings.write_time('GfasVerticalDistribution', 'calculate_widths', timeit.default_timer() - st_time, level=3)
        return widths

    def get_weights(self, heights_list):
        """
        Calculate the proportion (%) of emission to put on each layer.

        :param heights_list: List with the width of each vertical level.
        :type heights_list: list

        :return: List of the weight to apply to each layer.
        :rtype: list
        """
        st_time = timeit.default_timer()

        weights = []
        width_list = self.calculate_widths(heights_list)
        if self.approach == 'uniform':
            max_percent = 1.
        elif self.approach == '50_top':
            max_percent = 0.5
            width_list = width_list[0:-1]
        else:
            max_percent = 1.

        for width in width_list:
            weights.append((width * max_percent) / sum(width_list))
        if self.approach == '50_top':
            if len(heights_list) == 1:
                weights.append(1.)
            else:
                weights.append(0.5)

        settings.write_time('GfasVerticalDistribution', 'get_weights', timeit.default_timer() - st_time, level=3)
        return weights

    def apply_approach(self, top_fires):
        """
        Scatters the emissions vertically.

        :param top_fires: 4D array (time, level, latitude, longitude) with all the emission on each top layer.
        :type top_fires: numpy.array

        :return: 4D array (time, level, latitude, longitude) with all the emission distributed on all the involved
                layers.
        :rtype: numpy.array
        """
        import numpy as np

        st_time = timeit.default_timer()

        fires = np.zeros(top_fires.shape)
        for i in xrange(len(self.output_heights)):
            if top_fires[i].sum() != 0:
                weight_list = self.get_weights(list(self.output_heights[0: i + 1]))
                for i_weight in xrange(len(weight_list)):
                    fires[i_weight] += top_fires[i] * weight_list[i_weight]

        settings.write_time('GfasVerticalDistribution', 'apply_approach', timeit.default_timer() - st_time, level=3)
        return fires

    def do_vertical_interpolation_allocation(self, values, altitude):
        """
        Allocates the fire emissions on their top level.

        :param values: 2D array with the fire emissions
        :type values: numpy.array

        :param altitude: 2D array with the altitude of the fires.
        :type altitude: numpy.array

        :return: Emissions already allocated on the top altitude of each fire.
        :rtype: numpy.array
        """
        import numpy as np

        st_time = timeit.default_timer()

        fire_list = []
        aux_var = values
        for height in self.output_heights:
            aux_data = np.zeros(aux_var.shape)
            ma = np.ma.masked_less_equal(altitude, height)
            aux_data[ma.mask] += aux_var[ma.mask]
            aux_var -= aux_data
            fire_list.append(aux_data)
        fire_list = np.array(fire_list).reshape((len(fire_list), values.shape[1], values.shape[2]))

        settings.write_time('GfasVerticalDistribution', 'do_vertical_interpolation_allocation',
                            timeit.default_timer() - st_time, level=3)
        return fire_list

    def do_vertical_interpolation(self, values):
        """
        Manages all the process to do the vertical distribution.

        :param values: Emissions to be vertically distributed.
        :type values: numpy.array

        :return: Emissions already vertically distributed.
        :rtype: numpy.array
        """
        st_time = timeit.default_timer()

        fire_list = self.apply_approach(values)

        settings.write_time('GfasVerticalDistribution', 'do_vertical_interpolation', timeit.default_timer() - st_time,
                            level=3)
        return fire_list

    def calculate_weight_layer_dict(self, layer):
        weight_layer_dict = {x: None for x in xrange(layer + 1)}
        if self.approach == '50_top':
            weight_layer_dict[layer] = 0.5
            to_distribute = 0.5
            layer = layer - 1
        elif self.approach == 'uniform':
            to_distribute = 1.

        total_width = self.output_heights[layer]

        previous_height = 0
        for i, height in enumerate(self.output_heights[0:layer + 1]):
            partial_width = height - previous_height
            weight_layer_dict[i] = to_distribute * partial_width / total_width

            previous_height = height

        return weight_layer_dict

    def distribute_vertically(self, emissions_df, input_pollutant_list):
        import pandas as pd

        vert_emissions = []
        for layer, emis in emissions_df.groupby('layer'):
            if layer == 0:
                vert_emissions.append(emis)
            else:
                weight_layer_dict = self.calculate_weight_layer_dict(layer)
                for layer, weight in weight_layer_dict.iteritems():
                    aux_emis = emis.copy()
                    aux_emis.loc[:, 'layer'] = layer
                    aux_emis.loc[:, input_pollutant_list] = aux_emis[input_pollutant_list].multiply(weight)

                    vert_emissions.append(aux_emis)
        if len(vert_emissions) == 0:
            vert_emissions = emissions_df
        else:
            vert_emissions = pd.concat(vert_emissions)

        vert_emissions = vert_emissions.groupby(['FID', 'layer']).sum()
        vert_emissions.reset_index(inplace=True)
        return vert_emissions


if __name__ == '__main__':
    pass
