import datetime as dt
from dataclasses import replace

import matplotlib.pyplot as plt
import numpy as np
import pytest
import xarray as xr

from ravenpy.config.commands import BasinIndexCommand, HRUState
from ravenpy.config.rvs import RVC
from ravenpy.models import GR4JCN
from ravenpy.utilities.data_assimilation import assimilate, perturbation
from ravenpy.utilities.testdata import get_local_testdata


def test_perturbation():
    ts = get_local_testdata(
        "raven-gr4j-cemaneige/Salmon-River-Near-Prince-George_meteo_daily.nc"
    )
    ds = xr.open_dataset(ts)

    tmax = ds.tmax.isel(time=slice(0, 10))
    p_tmax = perturbation(tmax, "norm", 0.01, members=50)
    np.testing.assert_allclose(p_tmax.mean("members"), tmax, rtol=0.1)

    rain = ds.rain.isel(time=slice(30, 60))
    p_rain = perturbation(rain, "gamma", 0.01, members=50)
    np.testing.assert_allclose(p_rain.mean("members"), rain, rtol=0.1)

    assert p_tmax.attrs == ds.tmax.attrs
    assert p_rain.attrs == ds.rain.attrs


class TestAssimilationGR4JCN:
    def test_simple(self):

        # get timeseries
        ts = get_local_testdata(
            "raven-gr4j-cemaneige/Salmon-River-Near-Prince-George_meteo_daily.nc"
        )

        # set number of members. Using 7 here to make it easier to find and debug.
        n_members = 7

        # Perturbation parameters for the assimilation, keyed by standard_name
        std = {
            "rainfall": 0.30,
            "prsn": 0.30,
            "tasmin": 2.0,
            "tasmax": 2.0,
            "water_volume_transport_in_river_channel": 0.15,
        }

        # Use the same random seed for both tasmin and tasmax
        rs = np.random.SeedSequence(None).generate_state(1)[0]
        seed = {"tasmin": rs, "tasmax": rs}

        # Perturbation distribution
        dists = {
            "pr": "gamma",
            "rainfall": "gamma",
            "prsn": "gamma",
            "water_volume_transport_in_river_channel": "rnorm",
        }

        qkey = "water_volume_transport_in_river_channel"
        if qkey not in std:
            raise ValueError("Assimilation requires perturbing the flow variable.")

        # Assimilation variables (from HRUStateVariable)
        assim_var = ("soil0", "soil1")

        # Assimilation periods
        assim_days = [10] + 4 * [3]

        # GR4JCN model instance
        model = GR4JCN()

        # set the start and end dates for the first assimilation period, warm-up
        start_date = dt.datetime(2000, 6, 1)
        end_date = start_date + dt.timedelta(days=sum(assim_days))

        # Set model options

        model.config.rvh.hrus = (
            GR4JCN.LandHRU(
                area=4250.6, elevation=843.0, latitude=54.4848, longitude=-123.3659
            ),
        )

        model.config.rvp.params = GR4JCN.Params(
            0.1353389, -0.005067198, 576.8007, 6.986121, 1.102917, 0.9224778
        )  # SALMON

        # ==== Initialization (just to get reasonable states) ====
        # Set initialization run options
        model.config.rvi.run_name = "init"
        model.config.rvi.start_date = start_date
        # Richard: what is the end date policy for the init run ?
        model.config.rvi.end_date = start_date + dt.timedelta(days=assim_days[0])

        # Run the model
        model([ts])

        # Extract final model states
        hru_state, basin_state = model.get_final_state()
        xa = n_members * [getattr(hru_state, key) for key in assim_var]
        hru_states = n_members * [hru_state]
        basin_states = n_members * [basin_state]

        # === Create perturbed time series for full assimilation period ====
        perturbed = {}
        for key, s in std.items():
            nc = model.config.rvt._var_cmds[key]

            with xr.open_dataset(nc.file_name_nc) as ds:
                da = ds.get(nc.var_name_nc).sel(time=slice(start_date, end_date))

                perturbed[key] = perturbation(
                    da,
                    dists.get(key, "norm"),
                    std=s,
                    seed=seed.get(key, None),
                    member=n_members,
                )

                # Save flow for later
                if key == qkey:
                    q_obs = da

        # Write to disk
        p_fn = model.workdir / "perturbed_forcing.nc"
        perturbed = xr.Dataset(perturbed)
        perturbed.to_netcdf(p_fn, mode="w")

        # ==== Assimilation ====
        q_assim = []
        sd = start_date

        for i, ndays in enumerate(assim_days):

            dates = [sd + dt.timedelta(days=x) for x in range(ndays)]
            model.config.rvi.end_date = dates[-1]
            model.config.rvi.run_name = f"assim_{i}"

            # Perform the first assimilation step here
            [xa, model] = assimilate(
                model, p_fn, q_obs, assim_var, basin_states, hru_states, dates
            )

            # Save streamflow simulation
            q_assim.append(model.q_sim.isel(nbasins=0))

            # Update the start-time for the next loop
            sd += dt.timedelta(days=ndays)
            model.config.rvi.start_date = sd

            # Get new initial conditions and feed assimilated values
            hru_states, basin_states = model.get_final_state()
            hru_states = [
                replace(hru_states[i], **dict(zip(assim_var, xa[:, i])))
                for i in range(n_members)
            ]

        q_assim = xr.concat(q_assim, dim="time")

        # ==== Reference run ====
        model.config.rvi.run_name = "ref"
        model.config.rvi.start_date = start_date
        model.config.rvi.end_date = end_date

        model.config.rvc.hru_states = {}
        model.config.rvc.basin_states = {}
        model.config.rvc.soil0 = None
        model.config.rvc.soil1 = 15

        model([ts])

        # We can now plot everything!
        plt.plot(q_assim.T, "r", label="Assimilated")  # plot the assimilated flows
        plt.plot(q_obs.T, "b", label="Observed")  # plot the observed flows
        plt.plot(
            model.q_sim, "g", label="Simulated"
        )  # plot the open_loop (simulation with no assimilation)
        # plt.legend()
        # plt.show()

        assert q_assim.shape[0] == n_members
        assert q_assim.shape[1] == sum(assim_days)
