#!/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 parse_extrapolate(extrapolate) -> tuple:
    """
    Parses the 'extrapolate' parameter and returns a tuple representing the extrapolation options.

    Parameters
    ----------
    extrapolate : bool or tuple or None or number or NaN
        If bool:
            - If True, both extrapolation options are set to 'extrapolate'.
            - If False, extrapolation options are set to ('bottom', 'top').
        If tuple:
            - The first element represents the extrapolation option for the lower bound.
            - The second element represents the extrapolation option for the upper bound.
            - If any element is bool:
                - If True, it represents 'extrapolate'.
                - If False:
                    - If it's the first element, it represents 'bottom'.
                    - If it's the second element, it represents 'top'.
            - If any element is None, it is replaced with numpy.nan.
            - Other numeric values are kept as they are.
            - If any element is NaN, it is kept as NaN.
        If None:
            - Both extrapolation options are set to (NaN, NaN).
        If number:
            - Both extrapolation options are set to the provided number.
        If NaN:
            - Both extrapolation options are set to NaN.

    Returns
    -------
    tuple
        A tuple representing the extrapolation options. If the input is invalid, it returns
        ('extrapolate', 'extrapolate').
    """
    if isinstance(extrapolate, bool):
        if extrapolate:
            extrapolate_options = ('extrapolate', 'extrapolate')
        else:
            extrapolate_options = ('bottom', 'top')
    elif isinstance(extrapolate, tuple):
        extrapolate_options = [None, None]
        for i in range(len(extrapolate)):
            if isinstance(extrapolate[i], bool):
                if extrapolate[i]:
                    extrapolate_options[i] = 'extrapolate'
                else:
                    if i == 0:
                        extrapolate_options[i] = 'bottom'
                    else:
                        extrapolate_options[i] = 'top'
            elif extrapolate[i] is None:
                extrapolate_options[i] = np.nan
            else:
                extrapolate_options[i] = extrapolate[i]
        extrapolate_options = tuple(extrapolate_options)
    elif extrapolate is None:
        extrapolate_options = ('bottom', 'top')
    else:
        extrapolate_options = (extrapolate, extrapolate)

    return extrapolate_options


def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', extrapolate_options=False, info=None,
                         overwrite=False):
    """
    Vertical interpolation.

    Parameters
    ----------
    self : Nes
        Source Nes object.
    new_levels : List
        List of new vertical levels.
    new_src_vertical : 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.
    kind : str
        Vertical methods type.
    extrapolate_options : bool or tuple or None or number or NaN
        If bool:
            - If True, both extrapolation options are set to 'extrapolate'.
            - If False, extrapolation options are set to ('bottom', 'top').
        If tuple:
            - The first element represents the extrapolation option for the lower bound.
            - The second element represents the extrapolation option for the upper bound.
            - If any element is bool:
                - If True, it represents 'extrapolate'.
                - If False:
                    - If it's the first element, it represents 'bottom'.
                    - If it's the second element, it represents 'top'.
            - If any element is None, it is replaced with numpy.nan.
            - Other numeric values are kept as they are.
            - If any element is NaN, it is kept as NaN.
        If None:
            - Both extrapolation options are set to (NaN, NaN).
        If number:
            - Both extrapolation options are set to the provided number.
        If NaN:
            - Both extrapolation options are set to NaN.
    info: None, bool
        Indicates if you want to print extra information.
    overwrite: bool
        Indicates if you want to compute the vertical interpolation in the same object or not.
    """
    extrapolate_options = parse_extrapolate(extrapolate_options)
    do_extrapolation = 'extrapolate' in extrapolate_options

    if len(self.lev) == 1:
        raise RuntimeError("1D data cannot be vertically interpolated.")
    if not overwrite:
        self = self.copy(copy_vars=True)
    if info is None:
        info = self.info

    if new_src_vertical is not None:
        self.add_4d_vertical_info(new_src_vertical)
    if new_levels[0] > new_levels[-1]:
        ascendant = False
    else:
        ascendant = True

    nz_new = len(new_levels)

    if self.vertical_var_name is None:
        # To use current level data
        current_level = True
        # Checking old order
        src_levels = self.lev['data']
        if src_levels[0] > src_levels[-1]:
            if not ascendant:
                flip = False
            else:
                flip = True
                src_levels = np.flip(src_levels)
        else:
            if ascendant:
                flip = False
            else:
                flip = True
                src_levels = np.flip(src_levels)
    else:
        current_level = False
        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
        # Checking old order
        if np.nanmean(src_levels[:, 0, :, :]) > np.nanmean(src_levels[:, -1, :, :]):
            if not ascendant:
                flip = False
            else:
                flip = True
                src_levels = np.flip(src_levels, axis=1)
        else:
            if ascendant:
                flip = False
            else:
                flip = True
                src_levels = np.flip(src_levels, axis=1)

    # Loop over variables
    for var_name in self.variables.keys():
        if self.variables[var_name]['data'] is None:
            # Load data if it is not loaded yet
            self.load(var_name)

        if var_name != self.vertical_var_name:
            if flip:
                self.variables[var_name]['data'] = np.flip(self.variables[var_name]['data'], axis=1)
            if info and self.master:
                print("\t{var} vertical methods".format(var=var_name))
                sys.stdout.flush()
            nt, nz, ny, nx = self.variables[var_name]['data'].shape
            dst_data = np.empty((nt, nz_new, ny, nx), dtype=self.variables[var_name]['data'].dtype)
            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):
                        if len(src_levels.shape) == 1:
                            # To use 1D level information
                            curr_level_values = src_levels
                        else:
                            # To use 4D level data
                            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'

                            # Filtering filling values to extrapolation
                            fill_value = [np.nan, np.nan]
                            if 'bottom' in extrapolate_options:
                                if ascendant:
                                    fill_value[0] = np.float64(self.variables[var_name]['data'][t, 0, j, i])
                                else:
                                    fill_value[0] = np.float64(self.variables[var_name]['data'][t, -1, j, i])
                            else:
                                fill_value[0] = extrapolate_options[0]
                            if 'top' in extrapolate_options:
                                if ascendant:
                                    fill_value[1] = np.float64(self.variables[var_name]['data'][t, -1, j, i])
                                else:
                                    fill_value[1] = np.float64(self.variables[var_name]['data'][t, 0, j, i])
                            else:
                                fill_value[1] = extrapolate_options[1]
                            fill_value = tuple(fill_value)

                            # We force the methods with float64 to avoid negative values
                            # We don't know why the negatives appears with float34
                            if current_level:
                                # 1D vertical component
                                src_levels_aux = src_levels
                            else:
                                # 4D vertical component
                                src_levels_aux = src_levels[t, :, j, i]

                            if kind == 'linear' and ascendant and not do_extrapolation:
                                dst_data[t, :, j, i] = np.array(
                                    np.interp(new_levels,
                                              np.array(src_levels_aux, dtype=np.float64),
                                              np.array(self.variables[var_name]['data'][t, :, j, i], dtype=np.float64),
                                              left=fill_value[0], right=fill_value[1]),
                                    dtype=self.variables[var_name]['data'].dtype)
                            else:
                                if not do_extrapolation:
                                    dst_data[t, :, j, i] = np.array(
                                        interp1d(np.array(src_levels_aux, dtype=np.float64),
                                                 np.array(self.variables[var_name]['data'][t, :, j, i], dtype=np.float64),
                                                 kind=kind,
                                                 bounds_error=False,
                                                 fill_value=fill_value)(new_levels),
                                        dtype=self.variables[var_name]['data'].dtype)
                                else:
                                    # If extrapolation first we need to extrapolate all (below & above)
                                    dst_data[t, :, j, i] = np.array(
                                        interp1d(np.array(src_levels_aux, dtype=np.float64),
                                                 np.array(self.variables[var_name]['data'][t, :, j, i],
                                                          dtype=np.float64),
                                                 kind=kind,
                                                 bounds_error=False,
                                                 fill_value='extrapolate')(new_levels),
                                        dtype=self.variables[var_name]['data'].dtype)
                                    # Check values below the lower vertical level
                                    if fill_value[0] != 'extrapolate':
                                        if ascendant:
                                            idx_bellow = np.where(new_levels < src_levels_aux[0])
                                        else:
                                            idx_bellow = np.where(new_levels > src_levels_aux[0])
                                        dst_data[t, idx_bellow, j, i] = fill_value[0]
                                    # Check values above the upper vertical level
                                    if fill_value[1] != 'extrapolate':
                                        if ascendant:
                                            idx_above = np.where(new_levels > src_levels_aux[-1])
                                        else:
                                            idx_above = np.where(new_levels < src_levels_aux[-1])
                                        dst_data[t, idx_above, j, i] = fill_value[1]
                        except Exception as e:
                            print("time lat lon", t, j, i)
                            print("***********************")
                            print("LEVELS", src_levels_aux)
                            print("DATA", np.array(self.variables[var_name]['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'])

    # Update level information
    new_lev_info = {'data': np.array(new_levels)}
    if 'positive' in self._lev.keys():
        # Vertical level direction
        if flip:
            self.reverse_level_direction()
        new_lev_info['positive'] = self._lev['positive']

    if self.vertical_var_name is not None:
        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.free_vars(self.vertical_var_name)
        self.vertical_var_name = None

    self.set_levels(new_lev_info)

    # Remove original file information
    self.__ini_path = None
    self.dataset = None
    self.dataset = None

    return self
