#!/usr/bin/env python

import sys
import numpy as np
import os
import nes
from mpi4py import MPI
from scipy import spatial
from filelock import FileLock
from datetime import datetime
from warnings import warn


def interpolate_horizontal(self, dst_grid, weight_matrix_path=None, kind='NearestNeighbour', n_neighbours=4,
                           info=False):
    """
    Horizontal interpolation from one grid to another one.

    Parameters
    ----------
    self : nes.Nes
        Source projection Nes Object
    dst_grid : nes.Nes
        Final projection Nes object
    weight_matrix_path : str, None
        Path to the weight matrix to read/create
    kind : str
        Kind of horizontal interpolation. choices = ['NearestNeighbour']
    n_neighbours : int
        Used if kind == NearestNeighbour. Number of nearest neighbours to interpolate. default = 4
    info: bool
        Indicates if you want to print extra info during the interpolation process.
    """
    # Obtain weight matrix
    if self.parallel_method == 'T':
        weights, idx = get_weights_idx_t_axis(self, dst_grid, weight_matrix_path, kind, n_neighbours)
    elif self.parallel_method in ['Y', 'X']:
        weights, idx = get_weights_idx_xy_axis(self, dst_grid, weight_matrix_path, kind, n_neighbours)
    else:
        raise NotImplemented("Parallel method {0} is not implemented yet for horizontal interpolations. Use 'T'".format(
            self.parallel_method))

    # Apply weights
    final_dst = dst_grid.copy()
    final_dst.set_communicator(dst_grid.comm)
    # return final_dst
    final_dst.lev = self.lev
    final_dst._lev = self._lev
    final_dst.time = self.time
    final_dst._time = self._time
    final_dst.hours_start = self.hours_start
    final_dst.hours_end = self.hours_end

    for var_name, var_info in self.variables.items():
        if info and self.master:
            print("\t{var} horizontal interpolation".format(var=var_name))
            sys.stdout.flush()
        src_shape = var_info['data'].shape
        dst_shape = (src_shape[0], src_shape[1], idx.shape[-2], idx.shape[-1])
        # Creating new variable without data
        final_dst.variables[var_name] = {attr_name: attr_value for attr_name, attr_value in var_info.items()
                                         if attr_name != 'data'}
        # Creating empty data
        final_dst.variables[var_name]['data'] = np.empty(dst_shape)

        # src_data = var_info['data'].reshape((src_shape[0], src_shape[1], src_shape[2] * src_shape[3]))
        for time in range(dst_shape[0]):
            for lev in range(dst_shape[1]):
                src_aux = get_src_data(self.comm, var_info['data'][time, lev], idx, self.parallel_method)
                # src_aux = np.take(src_data[time, lev], idx)
                final_dst.variables[var_name]['data'][time, lev] = np.sum(weights * src_aux, axis=1)

    return final_dst


def get_src_data(comm, var_data, idx, parallel_method):
    """
    To obtain the needed src data to interpolate.

    Parameters
    ----------
    comm : MPI.Communicator

    var_data : np.array
        Rank source data.
    idx : np.array
        Index of the needed data in a 2D flatten way
    parallel_method: str
        Source parallel method

    Returns
    -------
    np.array
        Flatten source needed data
    """
    if parallel_method == 'T':
        var_data = var_data.flatten()
    else:
        var_data = comm.gather(var_data, root=0)
        if comm.Get_rank() == 0:
            if parallel_method == 'Y':
                axis = 0
            elif parallel_method == 'X':
                axis = 1
            else:
                raise NotImplementedError(parallel_method)
            var_data = np.concatenate(var_data, axis=axis)
            var_data = var_data.flatten()

        var_data = comm.bcast(var_data)

    var_data = np.take(var_data, idx)
    return var_data


# noinspection DuplicatedCode
def get_weights_idx_t_axis(self, dst_grid, weight_matrix_path, kind, n_neighbours):
    """
    To obtain the weights and source data index through the T axis.

    Parameters
    ----------
    self : nes.Nes
        Source projection Nes Object
    dst_grid : nes.Nes
        Final projection Nes object
    weight_matrix_path : str, None
        Path to the weight matrix to read/create
    kind : str
        Kind of horizontal interpolation. choices = ['NearestNeighbour']
    n_neighbours : int
        Used if kind == NearestNeighbour. Number of nearest neighbours to interpolate. default = 4

    Returns
    -------
    tuple
        weights and source data index

    """
    if weight_matrix_path is not None:
        with FileLock(weight_matrix_path.replace('.nc', '.lock')):
            if self.master:
                if os.path.isfile(weight_matrix_path):
                    weight_matrix = read_weight_matrix(weight_matrix_path, comm=MPI.COMM_SELF)
                    if len(weight_matrix.lev['data']) != n_neighbours:
                        warn("The selected weight matrix does not have the same number of nearest neighbours." +
                             "Re-calculating again but not saving it.")
                        if kind in ['NearestNeighbour', 'NearestNeighbours', 'nn', 'NN']:
                            weight_matrix = create_nn_weight_matrix(self, dst_grid, n_neighbours=n_neighbours)
                        else:
                            raise NotImplementedError(kind)
                else:
                    if kind in ['NearestNeighbour', 'NearestNeighbours', 'nn', 'NN']:
                        weight_matrix = create_nn_weight_matrix(self, dst_grid, n_neighbours=n_neighbours)
                    else:
                        raise NotImplementedError(kind)
                    if weight_matrix_path is not None:
                        weight_matrix.to_netcdf(weight_matrix_path)
            else:
                weight_matrix = None
    else:
        if kind in ['NearestNeighbour', 'NearestNeighbours', 'nn', 'NN']:
            if self.master:
                weight_matrix = create_nn_weight_matrix(self, dst_grid, n_neighbours=n_neighbours)
            else:
                weight_matrix = None
        else:
            raise NotImplementedError(kind)

    # Normalize to 1
    if self.master:
        weights = np.array(np.array(weight_matrix.variables['inverse_dists']['data'], dtype=np.float64) /
                           np.array(weight_matrix.variables['inverse_dists']['data'], dtype=np.float64).sum(axis=1),
                           dtype=np.float64)
        idx = np.array(weight_matrix.variables['idx']['data'][0], dtype=int)
    else:
        weights = None
        idx = None

    weights = self.comm.bcast(weights, root=0)
    idx = self.comm.bcast(idx, root=0)

    return weights, idx


# noinspection DuplicatedCode
def get_weights_idx_xy_axis(self, dst_grid, weight_matrix_path, kind, n_neighbours):
    """
    To obtain the weights and source data index through the X or Y axis.

    Parameters
    ----------
    self : nes.Nes
        Source projection Nes Object
    dst_grid : nes.Nes
        Final projection Nes object
    weight_matrix_path : str, None
        Path to the weight matrix to read/create
    kind : str
        Kind of horizontal interpolation. choices = ['NearestNeighbour']
    n_neighbours : int
        Used if kind == NearestNeighbour. Number of nearest neighbours to interpolate. default = 4

    Returns
    -------
    tuple
    weights and source data index

    """
    if weight_matrix_path is not None:
        with FileLock(weight_matrix_path.replace('.nc', '.lock')):
            if self.master:
                if os.path.isfile(weight_matrix_path):
                    weight_matrix = read_weight_matrix(weight_matrix_path, comm=MPI.COMM_SELF)
                    if len(weight_matrix.lev['data']) != n_neighbours:
                        warn("The selected weight matrix does not have the same number of nearest neighbours." +
                             "Re-calculating again but not saving it.")
                        if kind in ['NearestNeighbour', 'NearestNeighbours', 'nn', 'NN']:
                            weight_matrix = create_nn_weight_matrix(self, dst_grid, n_neighbours=n_neighbours)
                        else:
                            raise NotImplementedError(kind)
                else:
                    if kind in ['NearestNeighbour', 'NearestNeighbours', 'nn', 'NN']:
                        weight_matrix = create_nn_weight_matrix(self, dst_grid, n_neighbours=n_neighbours)
                    else:
                        raise NotImplementedError(kind)
                    if weight_matrix_path is not None:
                        weight_matrix.to_netcdf(weight_matrix_path)
            else:
                weight_matrix = None
    else:
        if kind in ['NearestNeighbour', 'NearestNeighbours', 'nn', 'NN']:
            if self.master:
                weight_matrix = create_nn_weight_matrix(self, dst_grid, n_neighbours=n_neighbours)
            else:
                weight_matrix = None
        else:
            raise NotImplementedError(kind)

    # Normalize to 1
    if self.master:
        weights = np.array(np.array(weight_matrix.variables['inverse_dists']['data'], dtype=np.float64) /
                           np.array(weight_matrix.variables['inverse_dists']['data'], dtype=np.float64).sum(axis=1),
                           dtype=np.float64)
        idx = np.array(weight_matrix.variables['idx']['data'][0], dtype=int)
    else:
        weights = None
        idx = None

    weights = self.comm.bcast(weights, root=0)
    weights = weights[:, :, dst_grid.write_axis_limits['y_min']:dst_grid.write_axis_limits['y_max'],
                      dst_grid.write_axis_limits['x_min']:dst_grid.write_axis_limits['x_max']]
    idx = self.comm.bcast(idx, root=0)
    idx = idx[:, dst_grid.write_axis_limits['y_min']:dst_grid.write_axis_limits['y_max'],
              dst_grid.write_axis_limits['x_min']:dst_grid.write_axis_limits['x_max']]

    return weights, idx


def read_weight_matrix(weight_matrix_path, comm=None, parallel_method='T'):
    """

    Parameters
    ----------
    weight_matrix_path : str
        Path of the weight matrix
    comm : MPI.Communicator
        Communicator to read the weight matrix
    parallel_method : str
        Nes parallel method to read the weight matrix.

    Returns
    -------
    nes.Nes
        Weight matrix
    """
    weight_matrix = nes.open_netcdf(path=weight_matrix_path, comm=comm, parallel_method=parallel_method, balanced=True)
    weight_matrix.load()

    return weight_matrix


def create_nn_weight_matrix(self, dst_grid, n_neighbours=4, info=False):
    """
    To create the weight matrix with the nearest neighbours method.

    Parameters
    ----------
    self : nes.Nes
        Source projection Nes Object
    dst_grid : nes.Nes
        Final projection Nes object
    n_neighbours : int
        Used if kind == NearestNeighbour. Number of nearest neighbours to interpolate. default = 4
    info: bool
        Indicates if you want to print extra info during the interpolation process.

    Returns
    -------
    nes.Nes
        Weight Matrix

    """
    if info and self.master:
        print("\tCreating Nearest Neighbour Weight Matrix with {0} neighbours".format(n_neighbours))
        sys.stdout.flush()
    # Source
    src_lat = np.array(self._lat['data'], dtype=np.float32)
    src_lon = np.array(self._lon['data'], dtype=np.float32)

    # 1D to 2D coordinates
    if len(src_lon.shape) == 1:
        src_lon, src_lat = np.meshgrid(src_lon, src_lat)

    # Destination
    dst_lat = np.array(dst_grid._lat['data'], dtype=np.float32)
    dst_lon = np.array(dst_grid._lon['data'], dtype=np.float32)

    # 1D to 2D coordinates
    if len(dst_lon.shape) == 1:
        dst_lon, dst_lat = np.meshgrid(dst_lon, dst_lat)

    # calculate N nearest neighbour inverse distance weights (and indices)
    # from gridcells centres of model 1 to each gridcell centre of model 2
    # model geographic longitude/latitude coordinates are first converted
    # to cartesian ECEF (Earth Centred, Earth Fixed) coordinates, before
    # calculating distances.

    src_mod_xy = lon_lat_to_cartesian(src_lon.flatten(), src_lat.flatten())
    dst_mod_xy = lon_lat_to_cartesian(dst_lon.flatten(), dst_lat.flatten())

    # generate KDtree using model 1 coordinates (i.e. the model grid you are
    # interpolating from)
    src_tree = spatial.cKDTree(src_mod_xy)

    # get n-neighbour nearest distances/indices (ravel form) of model 1 gridcell
    # centres from each model 2 gridcell centre

    dists, idx = src_tree.query(dst_mod_xy, k=n_neighbours)
    # self.nearest_neighbour_inds = \
    #     np.column_stack(np.unravel_index(idx, lon.shape))

    weight_matrix = dst_grid.copy()
    weight_matrix.time = [datetime(year=2000, month=1, day=1, hour=0, second=0, microsecond=0)]
    weight_matrix._time = [datetime(year=2000, month=1, day=1, hour=0, second=0, microsecond=0)]
    weight_matrix.set_communicator(MPI.COMM_SELF)
    # take the reciprocals of the nearest neighbours distances
    inverse_dists = np.reciprocal(dists)
    inverse_dists_transf = inverse_dists.T.reshape((1, n_neighbours, dst_lon.shape[0], dst_lon.shape[1]))
    weight_matrix.variables['inverse_dists'] = {'data': inverse_dists_transf, 'units': 'm'}
    idx_transf = idx.T.reshape((1, n_neighbours, dst_lon.shape[0], dst_lon.shape[1]))
    weight_matrix.variables['idx'] = {'data': idx_transf, 'units': ''}
    weight_matrix.lev = {'data': np.arange(inverse_dists_transf.shape[1]), 'units': ''}
    weight_matrix._lev = {'data': np.arange(inverse_dists_transf.shape[1]), 'units': ''}
    return weight_matrix


def lon_lat_to_cartesian(lon, lat, radius=1.0):
    """
    To calculate lon, lat coordinates of a point on a sphere

    Parameters
    ----------
    lon : np.array
        Longitude values
    lat : np.array
        Latitude values
    radius : float
        Radius of the sphere to get the distances.
    """
    lon_r = np.radians(lon)
    lat_r = np.radians(lat)

    x = radius * np.cos(lat_r) * np.cos(lon_r)
    y = radius * np.cos(lat_r) * np.sin(lon_r)
    z = radius * np.sin(lat_r)

    return np.column_stack([x, y, z])
