#!/usr/bin/env python3
"""Seed Default NAM Captures — lightweight synthetic captures for development.

Creates 3-5 default Neural Amp Modeler captures that cover the essential
guitar tones: clean, crunch, lead, rhythm, and hi-gain. Each capture uses
the Feather LSTM architecture (~350 parameters, runs in 0.5-3ms on RPi 4).

Usage
-----
    # Default: seed into ~/.pedal/nam/default/
    python3 scripts/seed_default_captures.py

    # Custom target directory
    python3 scripts/seed_default_captures.py --output /path/to/models

    # Dry run: print what would be created without writing
    python3 scripts/seed_default_captures.py --dry-run

    # As a module:
    from scripts.seed_default_captures import seed_captures
    seed_captures(output_dir="/home/pedal/models")

Output
------
    For each capture type, a ``.nam`` file is written containing JSON with
    architecture, config, and weights keys.

    A ``capture_index.json`` catalog is also written for UI discovery,
    conforming to the NAMModel metadata shape expected by
    ``NAMModel`` in ``src/dsp/nam_host.py``.

Dependencies
    - ``nam`` Python package (v0.13+)  — for model creation
    - Only requires stdlib + nam at runtime
"""

import argparse
import json
import os
import random
import sys
from pathlib import Path

try:
    import numpy as np
    import torch
    from nam.models import init_from_nam
except ImportError as exc:
    print(f"ERROR: {exc}. Install with: pip install nam", file=sys.stderr)
    sys.exit(1)

# ── Constants ────────────────────────────────────────────────────────────

DEFAULT_OUTPUT = Path.home() / ".pedal" / "nam" / "default"

CAPTURES = [
    {
        "id": "default_clean",
        "name": "Fender Twin Clean",
        "type": "clean",
        "category": "amp",
        "description": "Sparkling clean, slight headroom — edge of breakup when pushed",
        "gain_staging": 0.3,
        "seed": 42,
        "receptive_field": 33,
        "latent_size": 8,
        "sample_rate": 48000,
        "tags": ["clean", "fender", "twin", "default"],
    },
    {
        "id": "default_crunch",
        "name": "Marshall Plexi Crunch",
        "type": "crunch",
        "category": "amp",
        "description": "Classic rock crunch with smooth compression — Plexi on 7",
        "gain_staging": 0.8,
        "seed": 137,
        "receptive_field": 33,
        "latent_size": 8,
        "sample_rate": 48000,
        "tags": ["crunch", "marshall", "plexi", "rock", "default"],
    },
    {
        "id": "default_lead",
        "name": "Soldano Lead",
        "type": "lead",
        "category": "amp",
        "description": "Sustaining lead tone with fluid mids — singing harmonics",
        "gain_staging": 1.4,
        "seed": 256,
        "receptive_field": 65,
        "latent_size": 8,
        "sample_rate": 48000,
        "tags": ["lead", "soldano", "high-gain", "sustain", "default"],
    },
    {
        "id": "default_rhythm",
        "name": "Vox AC30 Rhythm",
        "type": "rhythm",
        "category": "amp",
        "description": "Chimey mid-gain rhythm — jangle with punch",
        "gain_staging": 0.6,
        "seed": 73,
        "receptive_field": 33,
        "latent_size": 8,
        "sample_rate": 48000,
        "tags": ["rhythm", "vox", "ac30", "jangle", "default"],
    },
    {
        "id": "default_high_gain",
        "name": "5150 Hi-Gain",
        "type": "hi-gain",
        "category": "amp",
        "description": "Modern high-gain with tight low-end — for metal and hard rock",
        "gain_staging": 2.0,
        "seed": 512,
        "receptive_field": 65,
        "latent_size": 12,
        "sample_rate": 48000,
        "tags": ["hi-gain", "5150", "metal", "modern", "default"],
    },
]


# ── Helpers ──────────────────────────────────────────────────────────────


def _lorenz(x0, y0, z0, sigma=10.0, rho=28.0, beta=8.0 / 3.0, steps=256):
    """Generate a chaotic attractor sequence (non-linear dynamics → amp-like)."""
    xs, ys, zs = [], [], []
    x, y, z = x0, y0, z0
    dt = 0.01
    for _ in range(steps):
        dx = sigma * (y - x)
        dy = x * (rho - z) - y
        dz = x * y - beta * z
        x += dx * dt
        y += dy * dt
        z += dz * dt
        xs.append(x)
        ys.append(y)
        zs.append(z)
    return np.array(ys)


def _generate_feather_weights(
    gain_staging: float,
    receptive_field: int,
    latent_size: int,
    seed: int,
    sample_rate: int = 48000,
) -> dict:
    """Generate synthetic Feather-style LSTM weights for a NAM capture.

    Uses deterministic chaotic dynamics (Lorenz attractor) seeded by the
    capture type to produce structured, non-random weight vectors. The
    result is a set of weights that a Feather LSTM model can load —
    different captures get different chaotic signatures → different
    tonal responses.

    Returns a dict matching the ``init_from_nam()`` format:
        {architecture, config, weights, sample_rate}
    """
    rng = random.Random(seed)
    torch.manual_seed(seed)

    # Feather LSTM: 1 LSTM layer, latent_size hidden, input_width=1 (audio)
    input_width = 1
    lstm_hidden = latent_size
    # output net: linear layer from latent_size -> 1

    # --- Compute exact parameter count for LSTM with this config ---
    # LSTM import_weights layout (per layer):
    #   Combined gate matrix: w_ih + w_hh = (4*hidden) * (input + hidden) elements
    #   bias_ih: 4 * hidden elements
    #   initial_hidden: hidden elements
    #   initial_cell: hidden elements
    #   head (Linear(hidden, 1)) weight: hidden elements + 1 bias = hidden + 1
    per_layer_gate = 4 * latent_size * (input_width + latent_size)
    per_layer_bias = 4 * latent_size
    per_layer_states = 2 * latent_size  # initial_hidden + initial_cell
    head_params = latent_size + 1  # Linear weight + bias

    total_params = per_layer_gate + per_layer_bias + per_layer_states + head_params

    # Generate weights using Lorenz attractor for structured non-linearity
    sig = _lorenz(
        x0=float(seed),
        y0=float(seed * 1.1),
        z0=float(seed * 0.9),
        steps=total_params + 500,
    )

    # Take the last `total_params` values and shape into weights
    sig = sig[-total_params:]

    # Apply gain staging + deterministic scaling
    sig = sig * gain_staging * 0.1

    # Normalize to prevent exploding activations
    sig = sig / (np.std(sig) + 1e-8)
    sig = np.clip(sig, -3.0, 3.0)

    weights = sig.tolist()

    # Build the config dict
    # Build the LSTM config dict — only keys that LSTM.__init__ accepts.
    # LSTM.__init__(self, hidden_size, input_size=1, sample_rate=None, **lstm_kwargs)
    # Extra keys go into **lstm_kwargs → nn.LSTM(...) and would crash.
    # Metadata (receptive_field, latent_size, type etc.) sit at TOP level,
    # outside "config", so init_from_nam never sees them.
    config = {
        "architecture": "LSTM",
        "config": {
            "hidden_size": latent_size,
            "input_size": input_width,
        },
        "weights": weights,
        "sample_rate": sample_rate,
    }

    return config


def _build_nam_file(config: dict) -> dict:
    """Test that ``init_from_nam()`` can load this config.

    Returns the config itself (pass-through). Raises if the model can't
    be instantiated.
    """
    # Validate that the model initializes
    try:
        model = init_from_nam(config)
        # Run dry inference to check shapes
        dummy = torch.randn(1, 256)  # one block of audio
        with torch.no_grad():
            out = model(dummy)
        # Note: actual model param count may exceed the flat weight array
        # because some params (bias_hh) are zeroed internally by import_weights
        # rather than loaded from the file. Only verify inference works.
        assert out.shape == dummy.shape, (
            f"Output shape {out.shape} != input shape {dummy.shape}"
        )
    except Exception as exc:
        raise RuntimeError(f"Failed to validate model: {exc}") from exc
    return config


def _write_capture(
    output_dir: Path,
    capture_def: dict,
    dry_run: bool = False,
) -> Path:
    """Generate and write one NAM capture file.

    Returns the path to the written file.
    """
    nam_config = _generate_feather_weights(
        gain_staging=capture_def["gain_staging"],
        receptive_field=capture_def["receptive_field"],
        latent_size=capture_def["latent_size"],
        seed=capture_def["seed"],
        sample_rate=capture_def["sample_rate"],
    )

    # Add metadata header for scan_models compatibility
    # The nam_host.py list_available_models() reads: architecture, config.receptive_field
    # We also embed UI-facing metadata
    nam_config["name"] = capture_def["name"]
    nam_config["type"] = capture_def["type"]
    nam_config["tags"] = capture_def["tags"]
    nam_config["description"] = capture_def["description"]
    nam_config["id"] = capture_def["id"]
    nam_config["category"] = capture_def["category"]

    # Validate the model can be built
    _build_nam_file(nam_config)

    output_path = output_dir / f"{capture_def['id']}.nam"
    if dry_run:
        param_count = len(nam_config["weights"])
        size_kb = param_count * 4 / 1024  # float32
        print(f"  [DRY-RUN] Would write: {output_path}")
        print(f"            Name: {capture_def['name']}")
        print(f"            Type: {capture_def['type']}")
        print(f"            Params: {param_count} ({size_kb:.1f} KB)")
        return output_path

    with open(output_path, "w") as f:
        json.dump(nam_config, f, indent=2)

    size_kb = output_path.stat().st_size / 1024
    print(f"  Wrote: {output_path.name} ({capture_def['name']}, {size_kb:.1f} KB)")
    return output_path


def _write_catalog(
    output_dir: Path,
    capture_paths: list[Path],
    capture_defs: list[dict],
    dry_run: bool = False,
) -> Path:
    """Write a ``capture_index.json`` catalog for UI consumption.

    This index mirrors the structure returned by
    ``NAMHost.list_available_models()`` in ``src/dsp/nam_host.py``
    plus additional UI metadata (type, tags, description).
    """
    entries = []
    for path, cdef in zip(capture_paths, capture_defs):
        size_bytes = 0 if dry_run else path.stat().st_size
        entries.append({
            "id": cdef["id"],
            "name": cdef["name"],
            "type": cdef["type"],
            "category": cdef["category"],
            "description": cdef["description"],
            "tags": cdef["tags"],
            "path": str(path),
            "size_bytes": size_bytes,
            "size_kb": round(size_bytes / 1024, 1),
            "architecture": "LSTM",
            "receptive_field": cdef["receptive_field"],
            "latent_size": cdef["latent_size"],
            "sample_rate": cdef["sample_rate"],
            "is_default": True,
        })

    catalog_path = output_dir / "capture_index.json"
    catalog = {
        "version": 1,
        "captures": entries,
    }

    if dry_run:
        print(f"  [DRY-RUN] Would write catalog: {catalog_path} ({len(entries)} entries)")
        return catalog_path

    with open(catalog_path, "w") as f:
        json.dump(catalog, f, indent=2)

    print(f"  Catalog: {catalog_path.name} ({len(entries)} entries, "
          f"{catalog_path.stat().st_size / 1024:.1f} KB)")
    return catalog_path


# ── Public API ───────────────────────────────────────────────────────────


def seed_captures(
    output_dir: str | Path = DEFAULT_OUTPUT,
    capture_types: list[str] | None = None,
    dry_run: bool = False,
) -> list[Path]:
    """Seed default NAM captures into ``output_dir``.

    Parameters
    ----------
    output_dir : str or Path
        Directory to write captures into. Created if it doesn't exist.
    capture_types : list of str or None
        Filter which captures to generate by type (e.g. ``["clean", "lead"]``).
        ``None`` generates all defined captures.
    dry_run : bool
        If True, print what would be done without writing files.

    Returns
    -------
    list[Path]
        Paths to the written ``.nam`` files.
    """
    output_dir = Path(output_dir)
    capture_defs = CAPTURES

    if capture_types is not None:
        capture_defs = [c for c in CAPTURES if c["type"] in capture_types]

    if not capture_defs:
        print(f"No captures match the requested types: {capture_types}")
        return []

    if dry_run:
        print(f"Dry-run: would seed {len(capture_defs)} captures to {output_dir}")
    else:
        output_dir.mkdir(parents=True, exist_ok=True)
        print(f"Seeding {len(capture_defs)} captures to {output_dir}")

    capture_paths = []
    for cdef in capture_defs:
        path = _write_capture(output_dir, cdef, dry_run=dry_run)
        capture_paths.append(path)

    _write_catalog(output_dir, capture_paths, capture_defs, dry_run=dry_run)

    total_params_approx = sum(
        (4 * c["latent_size"] * 1)  # input-hidden gates
        + (4 * c["latent_size"] * c["latent_size"])  # hidden-hidden gates
        + c["latent_size"]  # output projection
        for c in capture_defs
    )
    approx_size_kb = total_params_approx * 4 / 1024
    print(f"\nSummary: {len(capture_paths)} captures, ~{total_params_approx} total params, "
          f"~{approx_size_kb:.0f} KB total (uncompressed)")

    return capture_paths


# ── CLI ──────────────────────────────────────────────────────────────────


def main():
    parser = argparse.ArgumentParser(
        description="Seed default NAM captures for development and testing",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "--output", "-o",
        type=Path,
        default=DEFAULT_OUTPUT,
        help=f"Output directory (default: {DEFAULT_OUTPUT})",
    )
    parser.add_argument(
        "--types", "-t",
        nargs="+",
        choices=["clean", "crunch", "lead", "rhythm", "hi-gain"],
        help="Specific capture types to generate (default: all)",
    )
    parser.add_argument(
        "--dry-run", "-n",
        action="store_true",
        help="Print what would be created without writing files",
    )
    parser.add_argument(
        "--force", "-f",
        action="store_true",
        help="Overwrite existing .nam files",
    )

    args = parser.parse_args()

    # Check for existing captures
    if not args.dry_run and args.output.exists() and not args.force:
        existing = list(args.output.glob("default_*.nam"))
        if existing:
            print(
                f"Warning: {args.output} already contains {len(existing)} default "
                f"capture(s). Use --force to overwrite.",
                file=sys.stderr,
            )
            if not input("Continue anyway? [y/N] ").lower().startswith("y"):
                print("Aborted.")
                sys.exit(1)

    paths = seed_captures(
        output_dir=args.output,
        capture_types=args.types,
        dry_run=args.dry_run,
    )

    if not args.dry_run:
        print(f"\nDone. {len(paths)} captures written to {args.output}")
        print("To use: add a preset that references one of these .nam files, or")
        print("point NAMHost(models_dir=...) at this directory for scan_models().")


if __name__ == "__main__":
    main()