#!/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):
    """
    
    Parameters
    ----------
    self : nes.Nes
        
    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.

    Returns
    -------

    """
    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):
    """

    Parameters
    ----------
    self : Nes


    """
    if info is None:
        info = self.print_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 = src_data.shape
                dst_data = np.ma.masked_all((nt, nz_new, ny, nx))

                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(src_data[t, 0, j, i]), np.float64(src_data[t, -1, j, i]))
                            else:
                                fill_value = extrapolate
                            try:
                                # f = interp1d(np.array(self.lev['data'], dtype=np.float64),
                                #              np.array(src_data[t, :, j, i], dtype=np.float64),
                                #              kind=kind,
                                #              bounds_error=False,
                                #              fill_value=fill_value)
                                # dst_data[t, :, j, i] = np.array(f(new_levels), dtype=np.float32)
                                if kind == 'linear':
                                    dst_data[t, :, j, i] = np.array(
                                        np.interp(new_levels,
                                                  np.array(self.lev['data'], dtype=np.float64),
                                                  np.array(src_data[t, :, j, i], dtype=np.float64)), dtype=np.float32)
                                else:
                                    dst_data[t, :, j, i] = np.array(
                                        interp1d(np.array(self.lev['data'], dtype=np.float64),
                                                 np.array(src_data[t, :, j, i], dtype=np.float64),
                                                 kind=kind,
                                                 bounds_error=False,
                                                 fill_value=fill_value)(new_levels), dtype=np.float32)
                            except Exception as e:
                                print("time lat lon", t, j, i)
                                print("***********************")
                                print("LEVELS", self.lev['data'])
                                print("VAR", src_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.cumsum(np.flip(src_levels, axis=1), axis=1)
        else:
            src_levels = np.flip(src_levels, axis=1)
        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 = src_data.shape
                dst_data = np.empty((nt, nz_new, ny, nx), dtype=np.float32)

                for t in range(nt):
                    # if self.print_info and self.master:
                    if info and self.rank == self.size - 1:
                        print('\t\tTime {0} ({1}/{2}))'.format(self.time[t], t+1, nt))
                        sys.stdout.flush()
                    for j in range(ny):
                        for i in range(nx):
                            level_array = None
                            nl = src_levels[t, 0, j, i] - src_levels[t, 1, j, i]
                            if nl > 0:
                                level_max = np.max(src_levels[t, :, j, i])
                                level_array = np.asarray(new_levels)
                                level_array = level_array.astype('float32')
                                level_array = np.where(level_array > level_max, level_max, level_array)
                            if nl < 0:
                                level_min = np.min(src_levels[t, :, j, i])
                                level_array = np.asarray(new_levels)
                                level_array = level_array.astype('float32')
                                level_array = np.where(level_array < level_min, level_min, level_array)

                            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(src_data[t, 0, j, i]), np.float64(src_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(src_data[t, :, j, i], dtype=np.float64)), dtype=np.float32)
                                else:
                                    dst_data[t, :, j, i] = np.array(
                                        interp1d(np.array(src_levels[t, :, j, i], dtype=np.float64),
                                                 np.array(src_data[t, :, j, i], dtype=np.float64),
                                                 kind=kind,
                                                 bounds_error=False,
                                                 fill_value=fill_value)(new_levels), dtype=np.float32)
                            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(src_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 None
