#!/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 ESMF

import hermesv3_gr.config.settings as settings
from regrid import Regrid


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

        super(ConservativeRegrid, self).__init__(pollutant_dicts, weight_matrix_file, grid, masking=masking)

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

    def create_weight_matrix(self):
        """
        Calls to ESMF_RegridWeightGen to generate the weight matrix.
        """

        st_time = timeit.default_timer()

        src_grid = self.grid.create_esmf_grid_from_file(self.pollutant_dicts[0]['path'])
        src_field = ESMF.Field(src_grid, name='my input field')
        src_field.read(filename=self.pollutant_dicts[0]['path'], variable=self.pollutant_dicts[0]['name'],
                       timeslice=0)

        dst_grid = self.grid.esmf_grid
        dst_field = ESMF.Field(dst_grid, name='my outut field')
        ESMF.Regrid(src_field, dst_field, filename=self.weight_matrix_file, regrid_method=ESMF.RegridMethod.CONSERVE,)

        settings.write_time('ConservativeRegrid', 'create_weight_matrix', timeit.default_timer() - st_time, level=1)

    def start_regridding(self, gfas=False, vertical=None):
        # TODO Documentation
        from hermesv3_gr.tools.netcdf_tools import extract_vars

        st_time = timeit.default_timer()

        weights = self.read_weight_matrix()

        dst_field_list = []
        num = 1
        for pollutant_single_dict in self.pollutant_dicts:
            settings.write_log('\t\tPollutant {0} ({1}/{2})'.format(
                pollutant_single_dict['name'], num, len(self.pollutant_dicts)), level=3)
            num += 1

            [values] = extract_vars(pollutant_single_dict['path'], [pollutant_single_dict['name']])
            values = values['data']
            if gfas:
                values = vertical.do_vertical_interpolation_allocation(values, vertical.altitude)
            # 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)
            if gfas:
                values = values.reshape((values.shape[-3], values.shape[-2] * values.shape[-1],))
            else:
                values = values.reshape((1, values.shape[-2] * values.shape[-1],))

            unique, counts = np.unique(weights['row'], return_counts=True)
            new_dst_indices = np.cumsum(counts)

            # Expand src values
            src_aux = np.take(values, weights['col'], axis=1)

            # Apply weights
            dst_field_aux = np.multiply(src_aux, weights['S'])

            # Reduce dst values
            dst_field = self.reduce_dst_field(dst_field_aux, new_dst_indices, self.grid.shape[-1] * self.grid.shape[-2])

            if gfas:
                dst_field = vertical.do_vertical_interpolation(dst_field)
                dst_field = dst_field.reshape((self.grid.shape[-3], self.grid.shape[-2], self.grid.shape[-1],))
            else:
                dst_field = dst_field.reshape((self.grid.shape[-2], self.grid.shape[-1],))

            dst_field_list.append({'data': dst_field, 'name': pollutant_single_dict['name']})

        settings.write_time('ConservativeRegrid', 'start_regridding', timeit.default_timer() - st_time, level=3)
        return dst_field_list

    def read_weight_matrix(self):
        from netCDF4 import Dataset
        dict_aux = {}
        nc = Dataset(self.weight_matrix_file, mode='r')

        dict_aux['col'] = nc.variables['col'][:]
        dict_aux['row'] = nc.variables['row'][:]
        dict_aux['S'] = nc.variables['S'][:]
        nc.close()
        dict_aux['max'] = dict_aux['row'].max()

        dict_aux['col'] -= 1
        dict_aux['row'] -= 1

        if settings.size != 1:
            inc = dict_aux['row'][:-1] - dict_aux['row'][1:]
            index = np.where(inc > inc.max() * 0.5)[0]
            index = np.concatenate([[0], index])

            try:
                if settings.rank != 0:
                    dict_aux['col'] = dict_aux['col'][index[settings.rank] + 1: index[settings.rank + 1] + 1]
                    dict_aux['row'] = dict_aux['row'][index[settings.rank] + 1: index[settings.rank + 1] + 1]
                    dict_aux['S'] = dict_aux['S'][index[settings.rank] + 1: index[settings.rank + 1] + 1]
                else:
                    dict_aux['col'] = dict_aux['col'][: index[settings.rank + 1] + 1]
                    dict_aux['row'] = dict_aux['row'][: index[settings.rank + 1] + 1]
                    dict_aux['S'] = dict_aux['S'][: index[settings.rank + 1] + 1]
            except IndexError:
                dict_aux['col'] = dict_aux['col'][index[settings.rank] + 1:]
                dict_aux['row'] = dict_aux['row'][index[settings.rank] + 1:]
                dict_aux['S'] = dict_aux['S'][index[settings.rank] + 1:]

        return dict_aux

    def wait_to_weightmatrix(self):
        import time

        find = False

        while not find:
            if os.path.exists(self.weight_matrix_file):
                pre_size = 0
                post_size = 1
                print "I'm {0}".format(settings.rank), 'Writing Weight Matrix {0}'.format(self.weight_matrix_file)
                # find = True
                while pre_size != post_size:
                    print "I'm {0}".format(settings.rank), pre_size, post_size
                    pre_size = post_size
                    post_size = os.path.getsize(self.weight_matrix_file)
                    time.sleep(1)
                find = True
                print "I'm {0}".format(settings.rank), 'FINISHED'
            else:
                time.sleep(5)
                print "I'm {0}".format(settings.rank), 'Waiting Weight Matrix'

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

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

        :return: Values already regridded.
        :rtype: numpy.array
        """

        dst_field = super(ConservativeRegrid, self).apply_weights(values)

        return dst_field


if __name__ == '__main__':
    pass
