"""Fits simulation parameters to a target clip by CMA-ES over rendered batches.

    python3 ml/fit.py TARGET.wav --out FIT_DIR [--init E=1.7e11,h=8e-4,...] [--generations 12] [--popsize 16]

Each generation renders one batch of parameter sets with scripts/render-batch.ts
(one simulation per worker in parallel on the GPU), scores every clip against
the target, and updates the CMA-ES distribution. The free parameters, in a
normalized log (or linear) space, are

    E, h, Lx, Ly, sigma0, sigma1, strikeVelocity, strikeFracX, strikeFracY, bow sag (across the width)

with rho fixed at the default (E and rho enter the sound only through E / rho
at a fixed level, so one of them is redundant) and a single parabolic bow as
the shape. The objective combines the log-mel distance of match.py (timbre and
decay) with the fine-structure distances of the early and whole-clip spectra
(the positions of the partials). Every generation's best clip is kept in
FIT_DIR/best-GEN.wav and a listening page FIT_DIR/fit.html is written.
"""

import argparse
import json
import math
import os
import shutil
import subprocess
import sys
import time

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from match import read_wav_any, onset, segment, features  # noqa: E402
import linesim  # noqa: E402

ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
NODE = ["node", "--experimental-strip-types", os.path.join(ROOT, "scripts", "render-batch.ts")]

# name, kind (log or lin), lower, upper, default
SPACE = [
    ("E", "log", 60e9, 300e9, 170e9),
    ("h", "log", 0.3e-3, 2.0e-3, 0.79e-3),
    ("Lx", "log", 6 * 0.0254, 30 * 0.0254, 12 * 0.0254),
    ("Ly", "log", 6 * 0.0254, 30 * 0.0254, 18 * 0.0254),
    ("sigma0", "log", 0.15, 5.0, 1.15),
    ("sigma1", "log", 0.5e-4, 3e-3, 5e-4),
    ("strikeVelocity", "log", 0.8, 6.0, 3.0),
    ("strikeFracX", "lin", 0.5, 0.92, 0.7),
    ("strikeFracY", "lin", 0.12, 0.55, 0.35),
    ("bow", "lin", -6e-3, 6e-3, 3e-3),
    ("bowY", "lin", -6e-3, 6e-3, 0.0),
]
FREE = [n for n, *_ in SPACE]  # set from --free


class SimpleCMA:
    """Minimal (mu/mu_w, lambda)-CMA-ES (Hansen's tutorial, without the
    rank-mu weights beyond the elite) for when the cma package is absent.
    Candidates are clipped to the unit cube."""

    def __init__(self, x0, sigma, popsize, seed=1):
        self.n = len(x0)
        self.mean = np.array(x0, float)
        self.sigma = sigma
        self.lam = popsize
        self.mu = popsize // 2
        w = np.log(self.mu + 0.5) - np.log(np.arange(1, self.mu + 1))
        self.w = w / w.sum()
        self.mueff = 1 / np.sum(self.w ** 2)
        n = self.n
        self.cc = (4 + self.mueff / n) / (n + 4 + 2 * self.mueff / n)
        self.cs = (self.mueff + 2) / (n + self.mueff + 5)
        self.c1 = 2 / ((n + 1.3) ** 2 + self.mueff)
        self.cmu = min(1 - self.c1, 2 * (self.mueff - 2 + 1 / self.mueff) / ((n + 2) ** 2 + self.mueff))
        self.damps = 1 + 2 * max(0, math.sqrt((self.mueff - 1) / (n + 1)) - 1) + self.cs
        self.chiN = math.sqrt(n) * (1 - 1 / (4 * n) + 1 / (21 * n * n))
        self.pc = np.zeros(n); self.ps = np.zeros(n); self.C = np.eye(n)
        self.rng = np.random.default_rng(seed)
        self.gen = 0

    def ask(self):
        self.B, self.D = None, None
        vals, vecs = np.linalg.eigh(self.C)
        self.D = np.sqrt(np.maximum(vals, 1e-20)); self.B = vecs
        self.z = self.rng.standard_normal((self.lam, self.n))
        self.y = (self.B @ (self.D[:, None] * self.z.T)).T
        X = np.clip(self.mean + self.sigma * self.y, 0, 1)
        return [x for x in X]

    def tell(self, X, scores):
        order = np.argsort(scores)
        X = np.array(X)
        old = self.mean.copy()
        self.mean = self.w @ X[order[: self.mu]]
        y_w = (self.mean - old) / self.sigma
        C_inv_sqrt = self.B @ np.diag(1 / self.D) @ self.B.T
        self.ps = (1 - self.cs) * self.ps + math.sqrt(self.cs * (2 - self.cs) * self.mueff) * (C_inv_sqrt @ y_w)
        self.gen += 1
        hsig = np.linalg.norm(self.ps) / math.sqrt(1 - (1 - self.cs) ** (2 * self.gen)) / self.chiN < 1.4 + 2 / (self.n + 1)
        self.pc = (1 - self.cc) * self.pc + hsig * math.sqrt(self.cc * (2 - self.cc) * self.mueff) * y_w
        ys = (X[order[: self.mu]] - old) / self.sigma
        rank_mu = sum(wi * np.outer(yi, yi) for wi, yi in zip(self.w, ys))
        self.C = (1 - self.c1 - self.cmu) * self.C + self.c1 * (np.outer(self.pc, self.pc) + (1 - hsig) * self.cc * (2 - self.cc) * self.C) + self.cmu * rank_mu
        self.sigma *= math.exp((self.cs / self.damps) * (np.linalg.norm(self.ps) / self.chiN - 1))


FIXED = {}  # name -> value for the parameters that are not free


def free_space():
    return [e for e in SPACE if e[0] in FREE]


def to_unit(vals):
    """Physical values of the free parameters -> [0, 1] coordinates."""
    u = []
    for (name, kind, lo, hi, _), v in zip(free_space(), vals):
        u.append((math.log(v / lo) / math.log(hi / lo)) if kind == "log" else (v - lo) / (hi - lo))
    return np.array(u)


def from_unit(u):
    vals = []
    for (name, kind, lo, hi, _), x in zip(free_space(), np.clip(u, 0, 1)):
        vals.append(lo * (hi / lo) ** x if kind == "log" else lo + (hi - lo) * x)
    return vals


def params_of(vals):
    """Full parameter dict (free values plus the fixed ones) with the shape as a feature string."""
    d = {name: d0 for name, _, _, _, d0 in SPACE}
    d.update(FIXED)
    d.update({name: v for (name, *_), v in zip(free_space(), vals)})
    bow, bow_y = d.pop("bow"), d.pop("bowY")
    feats = [f"bow:x,{bow:.6g},0"] if abs(bow) > 1e-5 else []
    if abs(bow_y) > 1e-5:
        feats.append(f"bow:y,{bow_y:.6g},0")
    d["features"] = ";".join(feats) if feats else "none"
    return d


def objective(ref, seg, tol=None):
    """Default: the mixed log-mel and fine-structure objective. With tol set: the tolerant line measure of linesim.py."""
    if tol is not None:
        d, parts = linesim.distance(linesim.picture(ref, tol), linesim.picture(seg, tol))
        return d, parts
    fe = features(seg)
    d_mel = float(np.mean(np.abs(ref[3] - fe[3])))
    d_early = 1 - float(ref[0] @ fe[0])
    d_whole = 1 - float(ref[1] @ fe[1])
    return d_mel + 4 * d_early + 4 * d_whole, (d_mel, d_early, d_whole)


def render(jobs, out, duration, workers):
    os.makedirs(out, exist_ok=True)
    jf = os.path.join(out, "jobs.json")
    with open(jf, "w") as f:
        json.dump({"jobs": jobs}, f)
    env = dict(os.environ)
    env.pop("DISPLAY", None)
    subprocess.run(NODE + ["--batch", jf, "--out", out, "--workers", str(workers), "--duration", str(duration)], check=True, env=env, cwd=ROOT,
                   stdout=open(os.path.join(out, "render.log"), "a"), stderr=subprocess.STDOUT)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("target")
    ap.add_argument("--out", required=True)
    ap.add_argument("--init", default="", help="comma-separated name=value in SI for the starting point")
    ap.add_argument("--sigma", type=float, default=0.15, help="initial CMA-ES step in the unit cube")
    ap.add_argument("--generations", type=int, default=12)
    ap.add_argument("--popsize", type=int, default=16)
    ap.add_argument("--workers", type=int, default=16)
    ap.add_argument("--duration", type=float, default=3.0)
    ap.add_argument("--free", default=",".join(FREE), help="comma-separated names of the free parameters")
    ap.add_argument("--set", default="", help="comma-separated name=value (SI) for fixed parameters that differ from the defaults")
    ap.add_argument("--linesim", action="store_true", help="use the tolerant line measure as the objective")
    ap.add_argument("--tol-start", type=float, default=0.05)
    ap.add_argument("--tol-end", type=float, default=0.02)
    args = ap.parse_args()
    FREE[:] = [n.strip() for n in args.free.split(",") if n.strip()]
    for kv in args.set.split(","):
        if "=" in kv:
            k, v = kv.split("="); FIXED[k.strip()] = float(v)
    os.makedirs(args.out, exist_ok=True)
    x = read_wav_any(args.target)
    st = onset(x)
    ref_seg = segment(x, st)
    ref = ref_seg if args.linesim else features(ref_seg)
    init = {name: d for name, _, _, _, d in SPACE}
    init.update(FIXED)
    for kv in args.init.split(","):
        if "=" in kv:
            k, v = kv.split("=")
            init[k.strip()] = float(v)
    x0 = to_unit([init[name] for name, *_ in free_space()])
    print(f"target {args.target}: onset {st / 48000:.3f} s; start {params_of(from_unit(x0))}", flush=True)

    try:
        import cma
        es = cma.CMAEvolutionStrategy(x0, args.sigma, {"popsize": args.popsize, "bounds": [0, 1], "seed": 1, "verbose": -9})
    except ImportError:
        es = SimpleCMA(x0, args.sigma, args.popsize)
    log = open(os.path.join(args.out, "fit.jsonl"), "a")
    best_overall = (float("inf"), None, None)
    history = []
    t0 = time.time()
    for gen in range(args.generations):
        tol = args.tol_start + (args.tol_end - args.tol_start) * gen / max(1, args.generations - 1) if args.linesim else None
        X = es.ask()
        jobs = [{"id": f"g{gen:02d}-{k:02d}", "params": params_of(from_unit(u))} for k, u in enumerate(X)]
        gdir = os.path.join(args.out, f"gen{gen:02d}")
        render(jobs, gdir, args.duration, args.workers)
        scores = []
        for j, u in zip(jobs, X):
            p = os.path.join(gdir, "clips", j["id"] + ".wav")
            if not os.path.exists(p):
                scores.append(1e3); continue
            s, parts = objective(ref, segment(read_wav_any(p), 0), tol)
            scores.append(s)
            log.write(json.dumps({"gen": gen, "id": j["id"], "score": s, "tol": tol, "parts": parts, "params": j["params"]}) + "\n")
        log.flush()
        es.tell(X, scores)
        k = int(np.argmin(scores))
        best = (scores[k], jobs[k]["params"], os.path.join(gdir, "clips", jobs[k]["id"] + ".wav"))
        shutil.copy(best[2], os.path.join(args.out, f"best-{gen:02d}.wav"))
        if best[0] < best_overall[0]:
            best_overall = best
        history.append({"gen": gen, "best": best[0], "mean": float(np.mean(scores)), "params": best[1]})
        print(f"gen {gen}: best {best[0]:.3f} (mean {np.mean(scores):.3f}, overall {best_overall[0]:.3f}{f', tol {tol:.3f}' if tol else ''}) after {(time.time() - t0) / 60:.1f} min; "
              + " ".join(f"{n}={v:.4g}" for n, v in best[1].items() if n != "features") + f" {best[1]['features']}", flush=True)
        write_page(args, history, best_overall)
    print("best:", json.dumps(best_overall[1]))


def write_page(args, history, best_overall):
    rows = "".join(
        f"<tr><td>{h['gen']}</td><td>{h['best']:.3f}</td><td>{h['mean']:.3f}</td><td><audio controls preload='none' src='best-{h['gen']:02d}.wav'></audio></td>"
        + "".join(f"<td>{h['params'][n]:.4g}</td>" for n, *_ in SPACE if n not in ("bow", "bowY")) + f"<td class='l'>{h['params']['features']}</td></tr>" for h in history)
    html = f"""<!doctype html><meta charset="utf-8"><title>fit to {os.path.basename(args.target)}</title>
<style>body{{font:14px system-ui,sans-serif;margin:16px}} table{{border-collapse:collapse}} td,th{{padding:3px 8px;border-bottom:1px solid #ddd;text-align:right;white-space:nowrap}} td.l{{text-align:left}} audio{{height:28px;width:200px}}</style>
<h1>Simulation fitted to {os.path.basename(args.target)}</h1>
<p>Target: <audio controls src="target.wav"></audio> &nbsp; CMA-ES, population {args.popsize}, objective = log-mel distance + 4 (early fine-structure distance + whole-clip fine-structure distance). Best so far: {best_overall[0]:.3f}.</p>
<table><tr><th>gen</th><th>best</th><th>mean</th><th>best clip</th>{"".join(f"<th>{n}</th>" for n, *_ in SPACE if n not in ("bow", "bowY"))}<th class="l">shape</th></tr>{rows}</table>"""
    with open(os.path.join(args.out, "fit.html"), "w") as f:
        f.write(html)


if __name__ == "__main__":
    main()