"""
Created on Sun Nov  1 20:48:03 2020

@author: Richard
"""

from copy import deepcopy

import numpy as np
import xarray as xr

"""
model = Raven model instance, preset with parameters etc.
xa = the set of state variables. In this case, soil0 and soil1 from GR4JCN. Will eventually need to update for other models, other variables
ts = the timeseries of inputs needed by model (tas, pr, qobs, etc.)
days = data for the days of assimilation. Assimilation is performed after last day.
number_members = number of EnKF members
std = variables for uncertainty estimation of input and streamflow variables.
    ex: std = {"rainfall": 0.30,"prsn": 0.30, "tasmin": 2.0, "tasmax": 2.0, "water_volume_transport_in_river_channel": 0.15}
precip = standard deviation used to sample precip, uses gamma distribution (fraction of observed value)
temperature = standard deviation used to sample temperature, normal dist (degrees Celcius)
qobs = standard deviation used to sample observed streamflow, normal dist (fraction of observed value)
"""


def assimilate(model, ts, q_obs, keys, basin_states, hru_states, days):
    """Assimilate streamflow over one day.

    Parameters
    ----------
    model : raven.Model
      Raven model instance configured to run.
    ts : str, Path
      Perturbed time series.
    keys : tuple
      Name of hru_state attributes to be assimilated, for example ("soil0", "soil1").
    basin_states : sequence
      Model initial conditions, BasinStateVariables instances.
    hru_states : sequence
      Model initial conditions, HRUStateVariables instances.
    ts : str, Path, list
      Input netCDF file names.
    days : datetime
      Dates ...
    std : dict
      Standard deviation of the perturbation noise for each input variable, keyed by variable standard name.

    """
    qkey = "water_volume_transport_in_river_channel"

    if len(basin_states) != len(hru_states):
        raise ValueError("`basin_states` and `hru_states` must have the same length.")

    model = deepcopy(model)

    # Number of members
    n_members = len(basin_states)

    # Run simulation with perturbed inputs
    model(ts, hru_state=hru_states, basin_state=basin_states, nc_index=range(n_members))

    # Extract final states (n_states, n_members)
    f_hru_states, f_basin_states = model.get_final_state()
    x_matrix = np.array(
        [[getattr(state, key) for key in keys] for state in f_hru_states]
    ).T

    # Sanity check
    if x_matrix.shape != (len(keys), n_members):
        raise ValueError

    # vnames = [HRU_NC_MAP[k] for k in keys]
    # x_matrix = model.storage.isel(time=-1)[vnames].to_array()

    # Last time step for assimilation
    qsim_vector = model.q_sim.isel(time=-1, nbasins=0).values

    # If there are problems related to missing Qobs or other variables, do not assimilate.
    with xr.open_dataset(ts) as perturbed:
        if perturbed.isnull().any().to_array().any():
            xa = x_matrix
        else:
            # Prepare the perturbed Qobs and Qobs errors for each member
            qobs_pert = perturbed[qkey].sel(time=days[-1])
            qobs_error = q_obs.sel(time=days[-1]) - qobs_pert

            # Do the assimilation, return the assimilated states for the next period
            xa = update_state(
                x_matrix, qobs_pert.values, qobs_error.values, qsim_vector
            )

    return [xa, model]


def perturbation(da, dist, std, seed=None, **kwargs):
    """Return perturbed time series.

    Parameters
    ----------
    da : DataArray
      Input time series to be perturbed.
    dist : {"norm", "gamma", "rnorm"}
      Name of statistical distribution from which random perturbations are drawn. `rnorm` stands for relative normal, where the given standard deviation is multiplied by the value being perturbed.
    std : float
      Standard deviation of random perturbation.
    seed : int
      Seed for the random number generator. Setting the same seed for different variables will ensure the same
      perturbations are generated.
    kwargs : dict
      Name and size of additional dimensions, apart from time.

    """
    nt = len(da.time)
    size = list(kwargs.values()) + [nt]
    dims = list(kwargs.keys()) + ["time"]

    # Create random state object. If seed is None, the seed itself will be random.
    rs = np.random.RandomState(seed)

    if dist == "norm":
        r = rs.normal(0, std, size=size)
        out = da + xr.DataArray(r, dims=dims, coords={"time": da.time})

    elif dist == "rnorm":
        r = rs.normal(0, da * std, size=size)
        out = da + xr.DataArray(r, dims=dims, coords={"time": da.time})

    elif dist == "gamma":
        shape = (da ** 2) / (std * da) ** 2
        scale = ((std * da) ** 2) / da
        r = np.nan_to_num(rs.gamma(shape=shape, scale=scale, size=size), nan=0.0)
        out = xr.DataArray(r, dims=dims, coords={"time": da.time})

    else:
        raise AttributeError(f"{dist} is not supported.")

    out.attrs.update(da.attrs)
    return out


def update_state(x, qobs_pert, qobs_error, qsim):
    """
    Update model state by assimilation.

    Parameters
    ----------
    x : ndarray (n_states, n_members)
      Model state initial values.
    qobs_pert : ndarray (n_members)
      Perturbed observed streamflows.
    qobs_error : ndarray (n_members)
      Perturbation added to qobs to get qobs_pert.
    qsim : ndarray (n_members)
      Simulated streamflows.

    Returns
    -------
    ndarray (n_states, n_members)
      Model state values after assimilation.

    References
    ----------
    The Ensemble Kalman Filter: theoretical formulation and practical implementation, Evensen 2003
    https://link.springer.com/article/10.1007%2Fs10236-003-0036-9

    University of colorado report on the efficient implementation of EnKF, Jan Mandel, 2006
    http://ccm.ucdenver.edu/reports/rep231.pdf
    """
    n_states, n_members = np.shape(x)

    # Make sure arrays have shape (1, n_members)
    qobs_pert = np.atleast_2d(qobs_pert)
    qobs_error = np.atleast_2d(qobs_error)
    qsim = np.atleast_2d(qsim)

    z = np.dot(qsim, np.ones((n_members, 1)))
    ha = qsim - (z * np.ones((1, n_members))) / n_members
    y = qobs_pert - qsim

    # Equations 4.1 from Mandel, 2006
    re = np.dot(qobs_error, qobs_error.transpose()) / n_members
    p = re + (np.dot(ha, ha.transpose())) / (n_members - 1)
    m = np.dot(p ** -1, y)
    z = (ha.transpose()) * m
    a = (
        x
        - np.dot((np.dot(x, np.ones((n_members, 1)))), np.ones((1, n_members)))
        / n_members
    )
    xa = x + (np.dot(a, z)) / (n_members - 1)
    xa = np.maximum(xa, 0)

    return xa
