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


class Regrid(object):
    # TODO Documentation
    def __init__(self, pollutant_dicts, weight_matrix_file, grid, masking=None):
        st_time = timeit.default_timer()
        settings.write_log('\t\t\tInitializing Regrid.', level=3)

        self.grid = grid
        self.pollutant_dicts = pollutant_dicts
        self.weight_matrix_file = weight_matrix_file
        self.masking = masking

        if not self.is_created_weight_matrix(erase=False):
            settings.write_log("\t\t\tWeight matrix {0} is not created. ".format(weight_matrix_file) +
                               "Trying to create it", level=1)
            settings.comm.Barrier()
            self.create_weight_matrix()

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

    def create_weight_matrix(self):
        """
        This function is not used because all the child classes have to implement it.
        """
        pass
        # implemented on inner class

    def apply_weights(self, values):
        """
        Calculate the regridded values using the ESMF algorithm for a 3D array.

        :param values: Input values to regrid
        :type values: numpy.array

        :return: Values already regridded.
        :rtype: numpy.array
        """
        from netCDF4 import Dataset

        st_time = timeit.default_timer()

        # Read weight matrix
        nc_weights = Dataset(self.weight_matrix_file, mode='r')

        src_indices = nc_weights.variables['src_indices'][:]
        max_index = nc_weights.variables['dst_indices'][:].max() + 1
        dst_indices_counts = nc_weights.variables['dst_indices_count'][:]
        weights = nc_weights.variables['weights'][:]

        nc_weights.close()

        # Do masking
        if self.masking.regrid_mask is not None:
            values = np.multiply(values, self.masking.regrid_mask)
        # Do scalling
        if self.masking.scale_mask is not None:
            values = np.multiply(values, self.masking.scale_mask)
        values = values.reshape(values.shape[1], values.shape[2] * values.shape[3])

        # Expand src values
        src_aux = np.take(values, src_indices, axis=1)

        # Apply weights
        dst_field_aux = np.multiply(src_aux, weights)

        # Reduce dst values
        dst_field = self.reduce_dst_field(dst_field_aux, dst_indices_counts, max_index)

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

        return dst_field

    @staticmethod
    def reduce_dst_field(dst_field_extended, dst_indices, max_index):
        """
        Reduces the values of the regridded data.
        eg:
            dst_field_extended = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
                        goes to:  0  0  0  1  1  1  2  2  2  2

            dst_indices = [3, 6, 10]

            result = [0+1+2, 3+4+5+, 6+7+8+9]
            result = [3, 12, 30]

        :param dst_field_extended: Array with as many elements as interconnections between src and dst with the dst
        values to be gathered.
        :type dst_field_extended: numpy.array

        :param dst_indices: Array with the last element index to
        :type dst_indices: numpy.array

        :param max_index:
        :type max_index: int

        :return:
        :rtype: numpy.array
        """
        st_time = timeit.default_timer()

        # Create new
        dst_field = np.zeros((dst_field_extended.shape[0], max_index), dtype=settings.precision)
        # dst_field = np.zeros((dst_field_extended.shape[0], self.grid.shape[-1] * self.grid.shape[-2]))

        previous = 0
        count = 0
        for i in dst_indices:
            try:
                dst_field[:, count] = dst_field_extended[:, previous:i].sum(axis=1, dtype=settings.precision)
            except ValueError:
                pass
            count += 1
            previous = i

        settings.write_time('Regrid', 'reduce_dst_field', timeit.default_timer() - st_time, level=3)

        return dst_field

    def is_created_weight_matrix(self, erase=False):
        """
        Checks if the weight matrix is created

        :return: Boolean that indicates if the weight matrix is already created.
        :rtype: bool
        """
        if erase and settings.rank == 0:
            if os.path.exists(self.weight_matrix_file):
                os.remove(self.weight_matrix_file)

        return os.path.exists(self.weight_matrix_file)


if __name__ == '__main__':
    pass
