from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import torch

ROOT = Path(__file__).resolve().parents[1]
CORE = ROOT / "core"
if str(CORE) not in sys.path:
    sys.path.insert(0, str(CORE))

from lda.ks import KSParams, integrate_steps, random_initial_state  # noqa: E402
from lda.utils import cyclic_distances, ensemble_spread, gaspari_cohn, rmse  # noqa: E402


def make_params(
    n: int,
    length: float | None,
    delta3: float = 0.1,
    delta5: float = 0.1,
) -> KSParams:
    if length is None:
        length = 0.125 * n
    return KSParams(
        n=n,
        length=length,
        linear_growth=1.0,
        third_order_dispersion=delta3,
        fifth_order_dispersion=delta5,
        fourth_order_damping=1.0,
    )


def forecast_ensemble(
    ensemble: np.ndarray,
    n_steps: int,
    dt: float,
    params: KSParams,
) -> np.ndarray:
    return np.stack([integrate_steps(member, n_steps, dt, params) for member in ensemble])


def make_truth(params: KSParams, rng: np.random.Generator, spinup_steps: int, dt: float) -> np.ndarray:
    state = random_initial_state(params, rng, std=0.5, max_initial_mode=min(64, params.n // 4))
    return integrate_steps(state, spinup_steps, dt, params)


def make_initial_ensemble(
    truth: np.ndarray,
    params: KSParams,
    rng: np.random.Generator,
    n_members: int,
    perturbation_std: float,
) -> np.ndarray:
    members = []
    for _ in range(n_members):
        perturb = random_initial_state(
            params,
            rng,
            std=perturbation_std,
            max_initial_mode=min(64, params.n // 4),
        )
        members.append(truth + perturb)
    return np.asarray(members)


def inflate_ensemble(ensemble: np.ndarray, factor: float) -> np.ndarray:
    mean = ensemble.mean(axis=0, keepdims=True)
    return mean + factor * (ensemble - mean)


def physical_localization(obs_loc: int, n: int, radius: float) -> np.ndarray:
    distance = cyclic_distances(np.arange(n), float(obs_loc), n)
    return gaspari_cohn(distance, radius)


def spatial_latent_localization(
    obs_loc: int,
    n: int,
    latent_channels: int,
    kz: int,
    radius: float,
) -> np.ndarray:
    latent_spacing = n / float(kz)
    latent_positions = np.arange(kz) * latent_spacing
    distance = cyclic_distances(latent_positions, float(obs_loc), n)
    weights = gaspari_cohn(distance, radius)
    return np.tile(weights, latent_channels)


def ensrf_scalar_update(
    ensemble: np.ndarray,
    y_ens: np.ndarray,
    obs_value: float,
    obs_error_var: float,
    loc_weights: np.ndarray | None = None,
) -> np.ndarray:
    n_members = ensemble.shape[0]
    x_mean = ensemble.mean(axis=0)
    x_pert = ensemble - x_mean
    y_mean = float(y_ens.mean())
    y_pert = y_ens - y_mean
    var_y = float(y_pert @ y_pert / max(n_members - 1, 1))
    denom = var_y + obs_error_var
    if not np.isfinite(denom) or denom <= 0.0:
        return ensemble

    cov_xy = y_pert[:, None].T @ x_pert / max(n_members - 1, 1)
    cov_xy = cov_xy.ravel()
    if loc_weights is not None:
        cov_xy = cov_xy * loc_weights

    gain = cov_xy / denom
    alpha = 1.0 / (1.0 + np.sqrt(obs_error_var / denom))
    new_mean = x_mean + gain * (obs_value - y_mean)
    new_pert = x_pert - alpha * np.outer(y_pert, gain)
    return new_mean + new_pert


def physical_ensrf(
    ensemble: np.ndarray,
    obs_values: np.ndarray,
    obs_locs: np.ndarray,
    obs_error_std: float,
    loc_radius: float | None,
) -> np.ndarray:
    analysis = ensemble.copy()
    obs_error_var = obs_error_std**2
    for obs_value, obs_loc in zip(obs_values, obs_locs):
        y_ens = analysis[:, obs_loc]
        loc = None if loc_radius is None else physical_localization(int(obs_loc), analysis.shape[1], loc_radius)
        analysis = ensrf_scalar_update(analysis, y_ens, float(obs_value), obs_error_var, loc)
    return analysis


def load_checkpoint(model: torch.nn.Module, checkpoint: str | None) -> torch.nn.Module:
    if checkpoint:
        state = torch.load(checkpoint, map_location="cpu")
        state_dict = state.get("model_state_dict", state.get("state_dict", state))
        model.load_state_dict(state_dict, strict=False)
    model.eval()
    return model


def encode_states(model: torch.nn.Module, states: np.ndarray, mean: float, std: float) -> np.ndarray:
    x = torch.as_tensor(((states - mean) / std)[:, None, :], dtype=torch.float32)
    with torch.no_grad():
        z = model.encode(x).detach().cpu().numpy()
    return z


def decode_states(model: torch.nn.Module, latent: np.ndarray, mean: float, std: float) -> np.ndarray:
    z = torch.as_tensor(latent, dtype=torch.float32)
    with torch.no_grad():
        x = model.decode(z).detach().cpu().numpy()[:, 0, :]
    return x * std + mean


def global_latent_localized_ensrf(
    ensemble: np.ndarray,
    obs_values: np.ndarray,
    obs_locs: np.ndarray,
    obs_error_std: float,
    loc_radius: float,
    model: torch.nn.Module,
    norm_mean: float,
    norm_std: float,
) -> np.ndarray:
    physical = ensemble.copy()
    obs_error_var = obs_error_std**2
    for obs_value, obs_loc in zip(obs_values, obs_locs):
        z = encode_states(model, physical, norm_mean, norm_std)
        decoded = decode_states(model, z, norm_mean, norm_std)
        y_ens = decoded[:, obs_loc]
        z_flat = z.reshape(z.shape[0], -1)
        z_flat_a = ensrf_scalar_update(z_flat, y_ens, float(obs_value), obs_error_var, None)
        decoded_a = decode_states(model, z_flat_a, norm_mean, norm_std)
        loc = physical_localization(int(obs_loc), physical.shape[1], loc_radius)
        physical = physical + (decoded_a - physical) * loc[None, :]
    return physical


def spatial_latent_ensrf(
    ensemble: np.ndarray,
    obs_values: np.ndarray,
    obs_locs: np.ndarray,
    obs_error_std: float,
    loc_radius: float,
    model: torch.nn.Module,
    norm_mean: float,
    norm_std: float,
) -> np.ndarray:
    z = encode_states(model, ensemble, norm_mean, norm_std)
    n_members, latent_channels, kz = z.shape
    z_flat = z.reshape(n_members, -1)
    obs_error_var = obs_error_std**2
    for obs_value, obs_loc in zip(obs_values, obs_locs):
        decoded = decode_states(model, z_flat.reshape(n_members, latent_channels, kz), norm_mean, norm_std)
        y_ens = decoded[:, obs_loc]
        loc = spatial_latent_localization(int(obs_loc), ensemble.shape[1], latent_channels, kz, loc_radius)
        z_flat = ensrf_scalar_update(z_flat, y_ens, float(obs_value), obs_error_var, loc)
    return decode_states(model, z_flat.reshape(n_members, latent_channels, kz), norm_mean, norm_std)


def diagnostics(prior: np.ndarray, analysis: np.ndarray, truth: np.ndarray) -> tuple[float, float, float, float]:
    prior_mean = prior.mean(axis=0)
    analysis_mean = analysis.mean(axis=0)
    return (
        rmse(prior_mean, truth),
        rmse(analysis_mean, truth),
        ensemble_spread(prior),
        ensemble_spread(analysis),
    )
