#!/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):
    """

    Parameters
    ----------
    self : nes.Nes

    dst_grid : nes.Nes


    """
    # Obtain weight matrix
    weights, idx = get_weights_idx(self, dst_grid, weight_matrix_path, kind, n_neighbours)

    # Apply weights
    final_dst = dst_grid.copy()
    # 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 = 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_weights_idx(self, dst_grid, weight_matrix_path, kind, n_neighbours):
    if self.parallel_method != 'T':
        raise NotImplemented("Parallel method {0} is not implemented yet for horizontal interpolations. Use 'T'".format(
            self.parallel_method))
    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


def read_weight_matrix(weight_matrix_path, comm=None, parallel_method='T'):
    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):
    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, R=1):
    """
    calculates lon, lat coordinates of a point on a sphere with
    radius R
    """
    lon_r = np.radians(lon)
    lat_r = np.radians(lat)

    x = R * np.cos(lat_r) * np.cos(lon_r)
    y = R * np.cos(lat_r) * np.sin(lon_r)
    z = R * np.sin(lat_r)
    return np.column_stack([x, y, z])
