"""Runnable end-to-end demo with matplotlib snapshots.

Goal
----
Show a realistic usage loop:
  * initialize population via sampler
  * evaluate a noisy multi-objective toy function
  * update crits (raw floats)
  * refresh() -> losers
  * choose_parent() and perform_new(remove_looser=True)
  * save 2-panel snapshots via matplotlib:
      - left: objective space (mu) + semi-transparent uncertainty ellipses
      - right: rank-coordinates used by pareto_core (mode="ranks")

Run
---
  pip install -e ".[examples]"
  python examples/demo_noisy_zdt1_snapshots.py --out out/demo --epochs 30

Notes
-----
* Uses Agg backend to work headless (CI / servers).
* Criteria directions are "min" for both objectives.
"""

from __future__ import annotations

import argparse
import math
from pathlib import Path

import numpy as np

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
from matplotlib.colors import Normalize  # noqa: E402
from matplotlib.patches import Ellipse  # noqa: E402

from zen_fronts import ZenFronts
from zen_fronts.pareto_core import transform_to_ranks
from zen_fronts.selection import SelectionConfig
from zen_fronts.utils.boundary import clip01, reflect01


def sampler(rng: np.random.Generator) -> dict[str, float]:
    # 2D continuous parameter space: x1,x2 in [0,1]
    return {"x1": float(rng.random()), "x2": float(rng.random())}


def _to_logit(x: float, eps: float = 1e-6) -> float:
    x = float(np.clip(float(x), eps, 1.0 - eps))
    return float(math.log(x / (1.0 - x)))


def _from_logit(z: float) -> float:
    return float(1.0 / (1.0 + math.exp(-float(z))))


def mutator(parent_params: dict[str, float], rng: np.random.Generator, **kw):
    """Gaussian mutation in a bounded [0,1]^2 box.

    `boundary` modes:
      * clip    : cheap, but creates 'sticky' clusters at 0/1.
      * reflect : recommended default for demos.
      * logit   : smooth near boundaries, but a bit more math.
    """

    sigma = float(kw.get("sigma", 0.08))
    boundary = str(kw.get("boundary", "reflect"))

    if boundary == "logit":
        x1 = _from_logit(_to_logit(parent_params["x1"]) + float(rng.normal(0.0, sigma)))
        x2 = _from_logit(_to_logit(parent_params["x2"]) + float(rng.normal(0.0, sigma)))
    else:
        x1 = float(parent_params["x1"] + rng.normal(0.0, sigma))
        x2 = float(parent_params["x2"] + rng.normal(0.0, sigma))
        if boundary == "clip":
            x1 = clip01(x1)
            x2 = clip01(x2)
        else:
            # reflect (default)
            x1 = reflect01(x1)
            x2 = reflect01(x2)

    return {"x1": float(x1), "x2": float(x2)}, {"sigma": sigma, "boundary": boundary}


def noisy_zdt1_like(x1: float, x2: float, rng: np.random.Generator) -> tuple[float, float]:
    # A simple Pareto-like tradeoff. MINIMIZE both objectives.
    f1 = float(x1)
    g = float(1.0 + 9.0 * x2)
    f2 = float(g * (1.0 - np.sqrt(f1 / g)))

    # Add observational noise (to make sigma2 meaningful)
    f1 = float(f1 + rng.normal(0.0, 0.01))
    f2 = float(f2 + rng.normal(0.0, 0.02))
    return f1, f2


def _stable_norm_rank_of_value(col: np.ndarray, idx_i: int, v: float) -> float:
    """Normalized stable rank position in [0,1] for a hypothetical value.

    Mirrors `pareto_core.transform_to_ranks()` ordering for a single column:
    sort by (value asc, index asc) and return pos/(N-1).
    """

    col = np.asarray(col, dtype=np.float64)
    N = int(col.size)
    if N <= 1:
        return 0.0

    idx = np.arange(N, dtype=np.int64)
    less = int(np.sum(col < float(v)))
    # Stable tie-break (rare for floats, but keep semantics exact).
    tie = int(np.sum((col == float(v)) & (idx < int(idx_i))))
    pos = less + tie
    return float(pos / float(N - 1))


def snapshot(zf: ZenFronts, epoch: int, out_dir: Path, *, ellipse_sigma: float = 1.0) -> None:
    """Save a 2-panel snapshot.

    Left: objective space (mu) with semi-transparent 1σ ellipses from (sigma2_f1, sigma2_f2).
    Right: the same population, but in per-objective stable rank coordinates used by pareto_core
           when `mode="ranks"` (computed on *maximized* objectives).
    """

    pids = sorted(zf.points.keys())
    if not pids:
        return

    xs, ys, cs = [], [], []
    sx, sy = [], []
    for pid in pids:
        info = zf.info(pid)
        c1 = float(info["crits"]["f1"]["mu"])
        c2 = float(info["crits"]["f2"]["mu"])
        s1 = float(max(0.0, info["crits"]["f1"].get("sigma2", 0.0)))
        s2 = float(max(0.0, info["crits"]["f2"].get("sigma2", 0.0)))
        sel = info.get("selection")
        # Larger quality_score is better by default (quality_sign=-1).
        q = 0.0 if sel is None or sel.get("quality_score") is None else float(sel["quality_score"])

        xs.append(c1)
        ys.append(c2)
        cs.append(q)
        sx.append(math.sqrt(s1))
        sy.append(math.sqrt(s2))

    xs_a = np.asarray(xs, dtype=np.float64)
    ys_a = np.asarray(ys, dtype=np.float64)
    cs_a = np.asarray(cs, dtype=np.float64)
    sx_a = np.asarray(sx, dtype=np.float64)
    sy_a = np.asarray(sy, dtype=np.float64)

    vmin = float(np.min(cs_a))
    vmax = float(np.max(cs_a))
    if vmin == vmax:
        vmin -= 1.0
        vmax += 1.0
    norm = Normalize(vmin=vmin, vmax=vmax)

    fig, (ax_l, ax_r) = plt.subplots(1, 2, figsize=(12.4, 5.2))

    # ---- left: objective space (min/min) ----
    ax_l.scatter(xs_a, ys_a, c=cs_a, s=25, norm=norm)
    for x, y, dx, dy in zip(xs_a, ys_a, ellipse_sigma * sx_a, ellipse_sigma * sy_a):
        if dx <= 0.0 and dy <= 0.0:
            continue
        e = Ellipse((float(x), float(y)), width=float(2.0 * dx), height=float(2.0 * dy), angle=0.0)
        e.set_alpha(0.15)
        e.set_facecolor("0.2")
        e.set_edgecolor("none")
        ax_l.add_patch(e)

    ax_l.set_xlabel("f1 (mu, min)")
    ax_l.set_ylabel("f2 (mu, min)")
    ax_l.set_title("objective space (+ ellipses from sigma2)")

    # ---- right: rank coordinates of maximized objectives (mode='ranks') ----
    # pareto_core assumes maximization; CriteriaStore internally flips min->max by mu *= -1.
    X_max = np.column_stack((-xs_a, -ys_a))
    R = transform_to_ranks(X_max)

    ax_r.scatter(R[:, 0], R[:, 1], c=cs_a, s=25, norm=norm)

    col0 = X_max[:, 0]
    col1 = X_max[:, 1]
    N = int(X_max.shape[0])
    min_w = 0.01

    for i in range(N):
        v0 = float(col0[i])
        v1 = float(col1[i])
        dv0 = float(ellipse_sigma * sx_a[i])
        dv1 = float(ellipse_sigma * sy_a[i])

        if dv0 > 0.0:
            r0_lo = _stable_norm_rank_of_value(col0, i, v0 - dv0)
            r0_hi = _stable_norm_rank_of_value(col0, i, v0 + dv0)
            w0 = max(min_w, abs(r0_hi - r0_lo))
        else:
            w0 = min_w

        if dv1 > 0.0:
            r1_lo = _stable_norm_rank_of_value(col1, i, v1 - dv1)
            r1_hi = _stable_norm_rank_of_value(col1, i, v1 + dv1)
            w1 = max(min_w, abs(r1_hi - r1_lo))
        else:
            w1 = min_w

        e = Ellipse((float(R[i, 0]), float(R[i, 1])), width=float(w0), height=float(w1), angle=0.0)
        e.set_alpha(0.15)
        e.set_facecolor("0.2")
        e.set_edgecolor("none")
        ax_r.add_patch(e)

    ax_r.set_xlim(-0.02, 1.02)
    ax_r.set_ylim(-0.02, 1.02)
    ax_r.set_xlabel("rank(-f1_mu)  (pareto_core mode='ranks')")
    ax_r.set_ylabel("rank(-f2_mu)  (pareto_core mode='ranks')")
    ax_r.set_title("rank space (+ approx ellipses)")

    fig.suptitle(f"epoch={epoch} (color=quality_score)")
    fig.tight_layout()

    out_path = out_dir / f"snap_{epoch:03d}.png"
    fig.savefig(out_path, dpi=160)
    plt.close(fig)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=str, default="out/demo", help="Output dir for png snapshots")
    ap.add_argument("--epochs", type=int, default=20)
    ap.add_argument("--pop", type=int, default=40)
    ap.add_argument("--seed", type=int, default=123)
    ap.add_argument("--eval-seed", type=int, default=999)
    ap.add_argument("--n-samples", type=int, default=128)
    ap.add_argument("--percentile", type=float, default=0.25)
    ap.add_argument("--sigma", type=float, default=0.08)
    ap.add_argument(
        "--boundary",
        type=str,
        default="reflect",
        choices=("reflect", "clip", "logit"),
        help="How to keep params in [0,1] during mutation",
    )
    ap.add_argument("--ellipse-sigma", type=float, default=1.0, help="Ellipse radius in σ (default: 1.0)")
    ap.add_argument("--snapshot-every", type=int, default=4)
    args = ap.parse_args()

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    zf = ZenFronts(
        crits={"f1": "min", "f2": "min"},
        selection=SelectionConfig(
            n_samples=int(args.n_samples),
            percentile=float(args.percentile),
            seed=int(args.seed),
            quantiles_mode_i=2,
            collect_stats=True,
        ),
        sampler=sampler,
        mutator=mutator,
    )

    ids = zf.add_random_point(int(args.pop))

    eval_rng = np.random.default_rng(int(args.eval_seed))

    # Warm start: one evaluation per point so they're ready for refresh().
    for pid in ids:
        p = zf.params(pid)
        f1, f2 = noisy_zdt1_like(p["x1"], p["x2"], eval_rng)
        zf.update_crits(pid, {"f1": f1, "f2": f2}, t=0.0)

    for epoch in range(1, int(args.epochs) + 1):
        # Evaluate all active points once per epoch.
        for pid in list(zf.points.keys()):
            p = zf.params(pid)
            f1, f2 = noisy_zdt1_like(p["x1"], p["x2"], eval_rng)
            zf.update_crits(pid, {"f1": f1, "f2": f2}, t=float(epoch))

        losers = zf.refresh(now=float(epoch))

        if epoch == 1 or (int(args.snapshot_every) > 0 and epoch % int(args.snapshot_every) == 0) or epoch == int(args.epochs):
            snapshot(zf, epoch, out_dir, ellipse_sigma=float(args.ellipse_sigma))

        for loser_id in losers:
            parent_id = zf.choose_parent(loser_id)
            new_id, _ = zf.perform_new(
                parent_id,
                looser=loser_id,
                remove_looser=True,
                sigma=float(args.sigma),
                boundary=str(args.boundary),
            )
            # Evaluate the new child once, so it becomes ready for next refresh.
            p = zf.params(new_id)
            f1, f2 = noisy_zdt1_like(p["x1"], p["x2"], eval_rng)
            zf.update_crits(new_id, {"f1": f1, "f2": f2}, t=float(epoch) + 0.1)

    print(f"snapshots written to: {out_dir}")


if __name__ == "__main__":
    main()
