'''
Contains functions used for simulation, smoothing, interpolation, differentiation, and integration.
'''

from __future__ import annotations

from typing import Union

import numpy as np
from numpy.typing import ArrayLike, NDArray
import pandas as pd
from pint import Quantity
from scipy.integrate import cumulative_trapezoid
from scipy.interpolate import griddata, RegularGridInterpolator
from numba import njit

from .typedefs import Kwargs, MatchErrStr
from ._classvars import RAW_DF_COLUMNS, RAW_DF_DEFAULT_UNITS, PROCESSED_DF_COLUMNS
from . import _pint_utils as pu
from . import _smoothing

@njit(cache=True)
def sim_func(
    t: NDArray[np.float64],
    h: NDArray[np.float64],
    m_max: np.float64,
    slope: np.float64,
    bump_height: np.float64
    ) -> NDArray[np.float64]:
    '''
    Compute simulated data.

    Model function is a decreasing logistic function with a tiny Gaussian bump.

    Decreasing logistic function: L / (1 + e^[-k(t-t0)])

    Gaussian: A e^(-(t-mu)^2 / a^2)
    '''

    # choose minimum moment to be 1/len(h) of m_max
    m_maxes = np.linspace(m_max / len(h), m_max, len(h))

    # Choose x0 to be two-thirds of the temperature range
    t_range = t.max() - t.min()
    t0 = t.min() + (2./3.) * t_range

    # derivative at t0 is log_max * k / 4, so scale k inversely with t_range
    k = -4. / t_range * slope

    gaussian_amp = m_max * bump_height

    # choose the Gaussian spread to be about one-fifth of the temperature range
    gaussian_spr = t_range / 5.

    # make the Gaussian center vary with field
    mus = t.min() + np.linspace(t_range * (2./3.), t_range / 3., len(h))

    m_grid = np.zeros((len(h), len(t)), dtype=np.float64)

    for i in range(len(h)):
        m_grid[i, :] = (
            m_maxes[i] / (1. + np.exp(-k * (t-t0)))
            + gaussian_amp * np.exp(-((t-mus[i]) / gaussian_spr)**2)
        )

    return m_grid

def smooth_diff_int_processed(
    used_fields: NDArray[np.float64],
    df_arrays: tuple[NDArray[np.float64], ...],
    n: np.int64,
    tr: NDArray[np.float64],
    d: np.int64,
    lms: NDArray[np.float64],
    lm_g: np.float64,
    we: bool,
    me: Union[bool, ArrayLike, MatchErrStr],
    mk: Kwargs,
    az: bool,
    current_p_units: dict[str, str],
    sample_mass: Quantity
    ) -> pd.DataFrame:
    '''Perform the smoothing, differentiation, and integration on the processed data.'''

    r_units = RAW_DF_DEFAULT_UNITS
    p_cols = PROCESSED_DF_COLUMNS

    # smooth in the supplied units for T and M, respectively
    smooth_output = create_output_arr(used_fields, n, np.linspace(tr[0], tr[1], n))
    _smoothing.smooth_processed(used_fields, df_arrays, smooth_output, d, lms, lm_g, we, me, mk)

    # fill processed df with the smoothed data in the supplied units
    smooth_output = smooth_output.reshape((len(used_fields)+1) * n, 8)
    processed_df = pu.df_from_units(p_cols.values(), current_p_units.values(), smooth_output)

    # convert T, H, M, and M_err to the default raw data units
    pu.convert_processed_raw(
        processed_df,
        (p_cols['T'], p_cols['H'], p_cols['M'], p_cols['M_err']),
        (r_units['T'], r_units['H'], r_units['M'], r_units['M_err'])
    )

    current_p_units['T'] = r_units['T']
    current_p_units['H'] = r_units['H']
    current_p_units['M'] = r_units['M']
    current_p_units['M_err'] = r_units['M_err']

    # compute per-mass columns
    pu.create_per_mass_cols(processed_df, p_cols, r_units, sample_mass)

    # perform the differentiation and integration in the default raw data units
    diff_int_output = processed_df.pint.dequantify().values.reshape(len(used_fields)+1, n, 8)
    diff_int_processed(diff_int_output)
    diff_int_output = diff_int_output.reshape((len(used_fields)+1) * n, 8)

    # remove zeros at beginning if add_zeros is False
    if not az:
        diff_int_output = diff_int_output[n:, :]

    processed_df = pu.df_from_units(p_cols.values(), current_p_units.values(), diff_int_output)

    # convert dM_dT and Delta_SM, which are in units directly constructed from the
    # defaults for T, H, and M, to the default (standard cgs) raw data units.
    pu.convert_processed_raw(
        processed_df,
        (p_cols['dM_dT'], p_cols['Delta_SM']),
        (r_units['dM_dT'], r_units['Delta_SM'])
    )

    return processed_df

@njit(parallel=True, cache=True)
def create_output_arr(
    used_fields: NDArray[np.float64],
    n: np.int64,
    t_grid: NDArray[np.float64]
    ) -> NDArray[np.float64]:
    '''
    Return an output array to hold the processed data.

    Like a stack of df values: (0 + fields) x npoints x columns

    The prepended field row holds data for zero field, where moment is zero
    (for integration).

    The first depth layer, output[:, :, 0], holds the temperature grid for
    each field.

    The second depth layer excluding zero field, output[1:, :, 1], holds the
    field values.

    M_err and M_per_mass_err are NaN.
    '''

    nfields = len(used_fields)

    output = np.zeros((nfields + 1, n, 8))
    output[:, :, 0] = t_grid.repeat(nfields + 1).reshape((n, nfields + 1)).T
    output[1:, :, 1] = used_fields.repeat(n).reshape((nfields, n))
    output[:, :, 3] = np.nan
    output[:, :, 5] = np.nan

    return output

def diff_int_processed(output: NDArray[np.float64]) -> None:
    '''
    Perform the differentiation and integration for the processed data.

    Data indices corresponding to columns are
    T: 0, H: 1, M_per_mass: 4, dM_dT: 6, Delta_SM: 7

    output[0,:,0] is a 1D array of the temperature grid

    output[:,0,1] is a 1D array of the field values
    '''

    # gradient along the npoints temperatures (axis=1)
    output[1:, :, 6] = np.gradient(output[1:,:,4], output[0,:,0], axis=1, edge_order=2)

    # integrate derivatives along the fields (axis=0)
    output[:, :, 7] = cumulative_trapezoid(output[:,:,6], x=output[:,0,1], axis=0, initial=0.0)

def interp_diff_int_raw(
    raw_df: pd.DataFrame,
    n: np.int64,
    temp_diff: np.float64,
    fields: NDArray[np.float64],
    t_h_raw_conv_values: NDArray[np.float64]
    ) -> None:
    '''Perform the interpolation, differentiation, and integration on the raw data.'''

    t_lin, h_lin, m_per_mass_grid = interp_setup(raw_df, n, temp_diff, fields, t_h_raw_conv_values)

    dm_dt_call, delta_sm_call = interp_callables(t_lin, h_lin, m_per_mass_grid)

    raw_df[RAW_DF_COLUMNS['dM_dT']] = dm_dt_call(t_h_raw_conv_values)
    raw_df[RAW_DF_COLUMNS['Delta_SM']] = delta_sm_call(t_h_raw_conv_values)

def interp_setup(
    raw_df: pd.DataFrame,
    n: np.int64,
    temp_diff: np.float64,
    fields: NDArray[np.float64],
    t_h_raw_conv_values: NDArray[np.float64]
    ) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
    '''Return T and H linspaces, and M_per_mass grid.'''

    t_min, t_max = t_h_raw_conv_values[:, 0].min(), t_h_raw_conv_values[:, 0].max()
    num_temps = np.int64(np.around(n * (t_max - t_min) / temp_diff))

    t_lin = np.linspace(t_min, t_max, num_temps)
    h_lin = np.concatenate(([0.0], fields, [t_h_raw_conv_values[:, 1].max()]))

    t_grid, h_grid = np.meshgrid(t_lin, h_lin, indexing='ij')

    # interpolated moment grid
    m_per_mass_grid = griddata(
        (t_h_raw_conv_values[:, 0], t_h_raw_conv_values[:, 1]),
        raw_df[RAW_DF_COLUMNS['M_per_mass']],
        (t_grid, h_grid),
        method='linear'
    )

    # same, but with nearest method instead of linear
    m_per_mass_nearest = griddata(
        (t_h_raw_conv_values[:, 0], t_h_raw_conv_values[:, 1]),
        raw_df[RAW_DF_COLUMNS['M_per_mass']],
        (t_grid, h_grid),
        method='nearest'
    )

    # fill any missing values in the "linear" grid with those from the "nearest" grid
    m_per_mass_nan_idx = np.isnan(m_per_mass_grid)
    m_per_mass_grid[m_per_mass_nan_idx] = m_per_mass_nearest[m_per_mass_nan_idx]

    return t_lin, h_lin, m_per_mass_grid

def interp_callables(
    t_lin: NDArray[np.float64],
    h_lin: NDArray[np.float64],
    m_per_mass_grid: NDArray[np.float64]
    ) -> tuple[RegularGridInterpolator, RegularGridInterpolator]:
    '''
    Compute dM_dT and Delta_SM grids and return callable
    ``RegularGridInterpolator``\s for interpolating dM_dT and Delta_SM,
    respectively, back to original T and H points.
    '''

    # compute dM_dT and Delta_SM from the grid
    dm_dt_grid = np.gradient(m_per_mass_grid, t_lin, axis=0, edge_order=2)
    delta_sm_grid = cumulative_trapezoid(dm_dt_grid, x = h_lin, axis=1, initial=0.0)

    # interpolate dM_dT back to the original T, H points
    dm_dt_call = RegularGridInterpolator(
        (t_lin, h_lin),
        dm_dt_grid,
        method='nearest',
        bounds_error=False
    )

    # interpolate Delta_SM back to the original T, H points
    delta_sm_call = RegularGridInterpolator(
        (t_lin, h_lin),
        delta_sm_grid,
        method='nearest',
        bounds_error=False
    )

    return dm_dt_call, delta_sm_call

def bootstrap(
    used_fields: NDArray[np.float64],
    df_arrays: tuple[NDArray[np.float64], ...],
    n: np.int64,
    tr: NDArray[np.float64],
    d: np.int64,
    lms: NDArray[np.float64],
    we: bool,
    az: bool,
    current_p_units: dict[str, str],
    sample_mass: Quantity,
    n_bootstrap: np.int64,
    rng: np.random.Generator
    ) -> pd.DataFrame:
    '''Perform the bootstrap calculation on the processed data.'''

    r_units = RAW_DF_DEFAULT_UNITS
    p_cols = PROCESSED_DF_COLUMNS

    # bootstrap in the supplied units for T and M, respectively
    bootstrap_output = create_output_arr(used_fields, n, np.linspace(tr[0], tr[1], n))
    _smoothing.bootstrap(
        used_fields, df_arrays, bootstrap_output,
        d, lms, we, n_bootstrap, rng
    )

    # fill processed df with the smoothed data in the supplied units
    bootstrap_output = bootstrap_output.reshape((len(used_fields)+1) * n, 8)

    bootstrap_df = pu.df_from_units(p_cols.values(), current_p_units.values(), bootstrap_output)

    # convert T, H, M, and M_err to the default raw data units
    pu.convert_processed_raw(
        bootstrap_df,
        (p_cols['T'], p_cols['H'], p_cols['M'], p_cols['M_err']),
        (r_units['T'], r_units['H'], r_units['M'], r_units['M_err'])
    )

    current_p_units['T'] = r_units['T']
    current_p_units['H'] = r_units['H']
    current_p_units['M'] = r_units['M']
    current_p_units['M_err'] = r_units['M_err']

    # compute per-mass columns
    pu.create_per_mass_cols(bootstrap_df, p_cols, r_units, sample_mass)

    bootstrap_output = bootstrap_df.pint.dequantify().values

    # remove zeros at beginning if add_zeros is False
    if not az:
        bootstrap_output = bootstrap_output[n:, :]

    bootstrap_df = pu.df_from_units(p_cols.values(), current_p_units.values(), bootstrap_output)

    return bootstrap_df
