#!/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 timeit
import hermesv3_gr.config.settings as settings


class VerticalDistribution(object):
    """
    VerticalDistribution class that contains all the information to do the vertical distribution.

    :param vertical_id: ID of the vertical profile that appears in the vertical profile file.
    :type vertical_id: str

    :param vertical_profile_path: Path to the file that contains all the vertical profiles.
    :type vertical_profile_path: str

    :param vertical_output_profile: path to the file that contain the vertical description of the required output
    file.
    :type vertical_output_profile: str
    """
    def __init__(self, vertical_id, vertical_profile_path, vertical_output_profile):
        st_time = timeit.default_timer()
        settings.write_log('\t\tInitializing Vertical.', level=2)

        self.id = vertical_id

        self.output_heights = vertical_output_profile
        self.vertical_profile = self.get_vertical_profile(vertical_profile_path)

        settings.write_time('VerticalDistribution', 'Init', timeit.default_timer() - st_time, level=2)

    def get_vertical_profile(self, path):
        """
        Extract the vertical v_profile from the vertical v_profile file.

        :param path: Path to the file that contains all the vertical profiles.
        :type path: str

        :return: List of tuples of two values. Te first value of the tuple is the height of the layer and the second
            value is the quantity (%) of pollutant that goes into this layer.
        :rtype: list of tuple
        """
        import pandas as pd
        import re

        st_time = timeit.default_timer()
        settings.write_log("\t\t\tGetting vertical profile id '{0}' from {1} .".format(self.id, path), level=3)

        df = pd.read_csv(path, sep=';')
        try:
            v_profile = df.loc[df[df.ID == self.id].index[0]].to_dict()
        except IndexError:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: Vertical profile ID {0} is not in the {1} file.'.format(self.id, path))
            sys.exit(1)
        v_profile.pop('ID', None)
        v_profile['layers'] = list(map(int, re.split(', |,|; |;| ', v_profile['layers'])))
        v_profile['weights'] = list(map(float, re.split(', |,|; |;| ', v_profile['weights'])))

        if len(v_profile['layers']) != len(v_profile['weights']):
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError("ERROR: The number of layers and numbers os weight have to have the same length." +
                                     " The v_profile '{0}' of the '{1}' file doesn't match.".format(self.id, path))
            sys.exit(1)
        else:
            return_value = list(zip(v_profile['layers'], v_profile['weights']))

        settings.write_time('VerticalDistribution', 'get_vertical_profile', timeit.default_timer() - st_time, level=3)

        return return_value

    @staticmethod
    def get_vertical_output_profile(path):
        """
        Extract the vertical description of the desired output.

        :param path: Path to the file that contains the output vertical description.
        :type path: str

        :return: Heights of the output vertical layers.
        :rtype: list
        """
        import pandas as pd

        st_time = timeit.default_timer()
        settings.write_log('Calculating vertical levels from {0} .'.format(path))

        df = pd.read_csv(path)

        heights = df.height_magl.values

        settings.write_time('VerticalDistribution', 'get_vertical_output_profile', timeit.default_timer() - st_time,
                            level=3)

        return heights

    @staticmethod
    def get_weights(prev_layer, layer, in_weight, output_vertical_profile):
        """
        Calculate the weights for the given layer.

        :param prev_layer: Altitude of the low layer. 0 if it's the first.
        :type prev_layer: float

        :param layer: Altitude of the current layer.
        :type layer: float

        :param in_weight: Weights
        :param output_vertical_profile:
        :return:
        """
        st_time = timeit.default_timer()

        output_vertical_profile_aux = [s for s in output_vertical_profile if s >= prev_layer]
        output_vertical_profile_aux = [s for s in output_vertical_profile_aux if s < layer]

        output_vertical_profile_aux = [prev_layer] + output_vertical_profile_aux + [layer]

        index = len([s for s in output_vertical_profile if s < prev_layer])
        origin_diff_factor = in_weight / (layer - prev_layer)
        weight_list = []
        for i in range(len(output_vertical_profile_aux) - 1):
            weight = (abs(output_vertical_profile_aux[i] - output_vertical_profile_aux[i + 1])) * origin_diff_factor
            weight_list.append({'index': index, 'weight': weight})
            index += 1

        settings.write_time('VerticalDistribution', 'get_weights', timeit.default_timer() - st_time, level=3)

        return weight_list

    def calculate_weights(self):
        """
        Calculate the weights for all the vertical layers.

        :return: Weights that goes to each layer.
        :rtype: list of float
        """
        import numpy as np

        st_time = timeit.default_timer()
        settings.write_log("\t\tCalculating vertical weights.", level=3)

        weights = np.zeros(len(self.output_heights))
        prev_layer = 0
        for layer, weight in self.vertical_profile:
            if weight != float(0):
                for element in self.get_weights(prev_layer, layer, weight, self.output_heights):
                    weights[element['index']] += element['weight']

            prev_layer = layer

        settings.write_time('VerticalDistribution', 'calculate_weights', timeit.default_timer() - st_time, level=3)

        return weights

    @staticmethod
    def apply_weights(data, weights):
        """
        Calculate the vertical distribution using the given data and weights.

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

        :param weights: Weights of each layer.
        :type weights: numpy.array

        :return: Emissions already vertically distributed.
        :rtype: numpy.array
        """
        import numpy as np

        st_time = timeit.default_timer()

        data_aux = np.multiply(weights.reshape(weights.shape + (1, 1)), data)

        settings.write_time('VerticalDistribution', 'apply_weights', timeit.default_timer() - st_time, level=3)

        return data_aux

    @staticmethod
    def apply_weights_level(data, weight):
        st_time = timeit.default_timer()

        for emi in data:
            if emi['data'] is not 0:
                emi['data'] = emi['data'] * weight

        settings.write_time('VerticalDistribution', 'apply_weights_level', timeit.default_timer() - st_time, level=3)

        return data


if __name__ == '__main__':
    pass
