#!/usr/bin/env python

import sys
import nes
from scipy.interpolate import interp1d
import numpy as np
from copy import copy


def add_4d_vertical_info(self, info_to_add):
    """
    To add the vertical information from other source.
    
    Parameters
    ----------
    self : nes.Nes
        Source Nes object.
    info_to_add : nes.Nes, str
        Nes object with the vertical information as variable or str with the path to the NetCDF file that contains 
        the vertical data.
    """

    vertical_var = list(self.concatenate(info_to_add))
    self.vertical_var_name = vertical_var[0]

    return None


def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', extrapolate=None, info=None):
    """
    Vertical interpolation method.

    Parameters
    ----------
    self : Nes
        Source Nes object.
    new_levels : list
        List of new vertical levels.
    new_src_vertical
    kind : str
        Vertical interpolation type.
    extrapolate : None, tuple, str
        Extrapolate method (for non linear operations).
    info: None, bool
        Indicates if you want to print extra information.
    """
    
    if info is None:
        info = self.info

    if new_src_vertical is not None:
        self.add_4d_vertical_info(new_src_vertical)

    nz_new = len(new_levels)

    if self.vertical_var_name is None:
        # To use current level data
        for var_name, var_info in self.variables.items():
            if var_info['data'] is None:
                self.load(var_name)
            if var_name != self.vertical_var_name:
                if info and self.master:
                    print("\t{var} vertical interpolation".format(var=var_name))
                    sys.stdout.flush()
                # src_data = np.flip(var_info['data'], axis=1)
                nt, nz, ny, nx = var_info['data'].shape
                dst_data = np.ma.masked_all((nt, nz_new, ny, nx), dtype=var_info['data'].dtype)

                for t in range(nt):
                    if info and self.master:
                        print('\t\tTime step: {0} ({1}/{2}))'.format(self.time[t], t + 1, nt))
                        sys.stdout.flush()
                    for j in range(ny):
                        for i in range(nx):
                            if extrapolate is None:
                                fill_value = (np.float64(var_info['data'][t, 0, j, i]),
                                              np.float64(var_info['data'][t, -1, j, i]))
                            else:
                                fill_value = extrapolate
                            try:
                                if kind == 'linear':
                                    dst_data[t, :, j, i] = np.array(
                                        np.interp(new_levels,
                                                  np.array(self.lev['data'], dtype=np.float64),
                                                  np.array(var_info['data'][t, :, j, i], dtype=np.float64)),
                                        dtype=var_info['data'].dtype)
                                else:
                                    dst_data[t, :, j, i] = np.array(
                                        interp1d(np.array(self.lev['data'], dtype=np.float64),
                                                 np.array(var_info['data'][t, :, j, i], dtype=np.float64),
                                                 kind=kind,
                                                 bounds_error=False,
                                                 fill_value=fill_value)(new_levels),
                                        dtype=var_info['data'].dtype)
                            except Exception as e:
                                print("time lat lon", t, j, i)
                                print("***********************")
                                print("LEVELS", self.lev['data'])
                                print("VAR", var_info['data'][t, :, j, i])
                                print("+++++++++++++++++++++++")
                                raise Exception(str(e))
                self.variables[var_name]['data'] = copy(dst_data)
        # Updating level information
        self.lev['data'] = new_levels
        self._lev['data'] = new_levels
        # raise NotImplementedError('Vertical information with no 4D vertical data is not implemented')
    else:
        src_levels = self.variables[self.vertical_var_name]['data']
        if self.vertical_var_name == 'layer_thickness':
            src_levels = np.flip(np.cumsum(np.flip(src_levels, axis=1), axis=1))
        else:
            # src_levels = np.flip(src_levels, axis=1)
            pass
        for var_name, var_info in self.variables.items():
            if var_info['data'] is None:
                self.load(var_name)
            if var_name != self.vertical_var_name:
                if info and self.master:
                    print("\t{var} vertical interpolation".format(var=var_name))
                    sys.stdout.flush()
                # src_data = np.flip(var_info['data'], axis=1)
                nt, nz, ny, nx = var_info['data'].shape
                dst_data = np.empty((nt, nz_new, ny, nx), dtype=np.float32)
                for t in range(nt):
                    # if info and self.rank == self.size - 1:
                    if self.info and self.master:
                        print('\t\t{3} time step {0} ({1}/{2}))'.format(self.time[t], t+1, nt, var_name))
                        sys.stdout.flush()
                    for j in range(ny):
                        for i in range(nx):
                            curr_level_values = src_levels[t, :, j, i]
                            try:
                                # check if all values are identical or masked
                                if ((isinstance(curr_level_values, np.ndarray) and
                                     (curr_level_values == curr_level_values[0]).all()) or
                                        (isinstance(curr_level_values, np.ma.core.MaskedArray) and
                                         curr_level_values.mask.all())):
                                    kind = 'slinear'
                                else:
                                    kind = kind  # 'cubic'
                                if extrapolate is None:
                                    fill_value = (np.float64(var_info['data'][t, 0, j, i]),
                                                  np.float64(var_info['data'][t, -1, j, i]))
                                else:
                                    fill_value = extrapolate

                                if kind == 'linear':
                                    dst_data[t, :, j, i] = np.array(
                                        np.interp(new_levels,
                                                  np.array(src_levels[t, :, j, i], dtype=np.float64),
                                                  np.array(var_info['data'][t, :, j, i], dtype=np.float64)),
                                        dtype=var_info['data'].dtype)
                                else:
                                    dst_data[t, :, j, i] = np.array(
                                        interp1d(np.array(src_levels[t, :, j, i], dtype=np.float64),
                                                 np.array(var_info['data'][t, :, j, i], dtype=np.float64),
                                                 kind=kind,
                                                 bounds_error=False,
                                                 fill_value=fill_value)(new_levels),
                                        dtype=var_info['data'].dtype)
                            except Exception as e:
                                print("time lat lon", t, j, i)
                                print("***********************")
                                print("LEVELS", np.array(src_levels[t, :, j, i], dtype=np.float64))
                                print("DATA", np.array(var_info['data'][t, :, j, i], dtype=np.float64))
                                print("METHOD", kind)
                                print("FILL_VALUE", fill_value)
                                print("+++++++++++++++++++++++")
                                raise Exception(str(e))
                            # if level_array is not None:
                            #     dst_data[t, :, j, i] = np.array(f(level_array), dtype=np.float32)

                self.variables[var_name]['data'] = copy(dst_data)
                # print(self.variables[var_name]['data'])

        # Updating level information
        new_lev_info = {'data': new_levels}
        for var_attr, attr_info in self.variables[self.vertical_var_name].items():
            if var_attr not in ['data', 'dimensions', 'crs', 'grid_mapping']:
                new_lev_info[var_attr] = copy(attr_info)
        self.lev = new_lev_info
        self._lev = new_lev_info
        self.free_vars(self.vertical_var_name)
        self.vertical_var_name = None

    return self
