from __future__ import annotations

import numpy as np
import matplotlib.pyplot as plt

from forcing import lorenz63_regime_switch, Lorenz63Params


# ============================================================
# Utilities
# ============================================================

def standardize_train_only(F_train: np.ndarray, F_test: np.ndarray, eps: float = 1e-12):
    mu = F_train.mean(axis=0, keepdims=True)
    sd = F_train.std(axis=0, keepdims=True) + eps
    return (F_train - mu) / sd, (F_test - mu) / sd, mu, sd


def make_task(F_test_z: np.ndarray, a_start: int, ell_ctx: int, steps: int):
    """
    prefix_z: (ell_ctx,3)
    truth_z : (steps,3)
    """
    prefix = F_test_z[a_start - ell_ctx : a_start, :]
    truth  = F_test_z[a_start : a_start + steps, :]
    return prefix, truth


def decimate(X: np.ndarray, max_points: int = 8000) -> np.ndarray:
    X = np.asarray(X)
    n = X.shape[0]
    if n <= max_points:
        return X
    stride = max(1, n // max_points)
    return X[::stride]


def set_equal_3d_limits(ax, all_xyz: np.ndarray, pad: float = 0.05):
    mins = all_xyz.min(axis=0)
    maxs = all_xyz.max(axis=0)
    ctr = 0.5 * (mins + maxs)
    span = float(np.max(maxs - mins)) * (1.0 + pad)

    ax.set_xlim(ctr[0] - span/2, ctr[0] + span/2)
    ax.set_ylim(ctr[1] - span/2, ctr[1] + span/2)
    ax.set_zlim(ctr[2] - span/2, ctr[2] + span/2)


def plot_context_grid_3d(
    traj_B_u: np.ndarray,
    preds_lisa_u: dict[int, np.ndarray],
    preds_alsa_u: dict[int, np.ndarray],
    *,
    title: str,
    elev: int = 20,
    azim: int = -60,
    max_bg_points: int = 12000,
):
    """
    2 x M grid:
      row 0: LISA predictions
      row 1: ALSA predictions
    Each panel has regime B attractor as gray background.
    """
    multipliers = sorted(preds_lisa_u.keys())
    M = len(multipliers)

    B = decimate(traj_B_u, max_bg_points)

    # Global axis limits based on everything (background + all predictions)
    all_parts = [B]
    for m in multipliers:
        all_parts.append(preds_lisa_u[m])
        all_parts.append(preds_alsa_u[m])
    all_xyz = np.concatenate(all_parts, axis=0)

    fig = plt.figure(figsize=(4.0 * M, 9.0))
    fig.suptitle(title, fontsize=14)

    # ---- helper panel ----
    def panel(ax, pred_u: np.ndarray, name: str, color: str):
        # regime B attractor background
        ax.plot(B[:, 0], B[:, 1], B[:, 2], color="0.80", lw=0.6, alpha=0.55)

        # prediction
        P = np.asarray(pred_u)
        ax.plot(P[:, 0], P[:, 1], P[:, 2], color=color, lw=2.2, alpha=0.95)

        # start/end markers
        ax.scatter(P[0, 0], P[0, 1], P[0, 2], s=25, color=color, alpha=0.95, marker="o")
        ax.scatter(P[-1, 0], P[-1, 1], P[-1, 2], s=25, color=color, alpha=0.95, marker="^")

        ax.set_title(name)
        ax.set_xlabel("x"); ax.set_ylabel("y"); ax.set_zlabel("z")
        ax.view_init(elev=elev, azim=azim)
        set_equal_3d_limits(ax, all_xyz)

        # cleaner look
        ax.grid(False)
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_zticks([])

    # Row 1: LISA
    for j, m in enumerate(multipliers):
        ax = fig.add_subplot(2, M, 1 + j, projection="3d")
        panel(ax, preds_lisa_u[m], f"LISA  ℓ={m}L", color="C1")

    # Row 2: ALSA
    for j, m in enumerate(multipliers):
        ax = fig.add_subplot(2, M, M + 1 + j, projection="3d")
        panel(ax, preds_alsa_u[m], f"ALSA  ℓ={m}L", color="C3")

    plt.tight_layout()
    plt.savefig("forcing_test.png")
    plt.show()
    return None

def unstandardize(Z: np.ndarray, sd) -> np.ndarray:
    return Z * sd + mu