"""Matched full-pipeline optimization trials, with load and warm runs separate."""

from __future__ import annotations

import argparse
import gc
import json
import logging
from pathlib import Path

from h3_mlx.runtime import Pipeline
from h3_mlx.weights import Linear


class HalfProjection(Linear):
    """Experimental FP16 GEMM, with BF16 rounding at every native boundary."""

    def __call__(self, x):
        import mlx.core as mx

        return super().__call__(x).astype(mx.bfloat16)


PROFILES = {
    "quiet-baseline": {"wired_gb": 42, "cache_refined_text": False},
    "quiet-native": {
        "wired_gb": 42,
        "compile_blocks": True,
        "eval_every": 0,
        "prefetch_blocks": 4,
        "vae_batch": 4,
        "vae_compile": True,
    },
    "quiet-packed": {
        "wired_gb": 32,
        "compile_blocks": True,
        "eval_every": 0,
        "prefetch_blocks": 4,
        "weight_storage": "packed12",
        "vae_batch": 4,
        "vae_compile": True,
    },
    "speed": {
        "wired_gb": 42,
        "compile_blocks": True,
        "eval_every": 0,
        "vae_batch": 4,
        "vae_compile": True,
    },
    "prefetch": {
        "wired_gb": 42,
        "compile_blocks": True,
        "eval_every": 0,
        "prefetch_blocks": 4,
        "vae_batch": 4,
        "vae_compile": True,
    },
    "compact": {
        "wired_gb": 32,
        "compile_blocks": True,
        "eval_every": 0,
        "prefetch_blocks": 4,
        "weight_storage": "packed12",
        "vae_batch": 4,
        "vae_compile": True,
    },
    "baseline": {},
    "compiled": {"compile_blocks": True},
    "compiled-batched": {
        "compile_blocks": True,
        "eval_every": 0,
        "keep_refiners": True,
    },
    "compiled-fast": {
        "compile_blocks": True,
        "eval_every": 0,
        "keep_refiners": True,
        "vae_batch": 8,
    },
    "fused-compiled": {
        "compile_blocks": True,
        "eval_every": 0,
        "keep_refiners": True,
        "vae_batch": 8,
        "kernel": "metal-qk",
    },
    "batched": {"eval_every": 0, "keep_refiners": True},
    "wired-baseline": {"wired_gb": 42},
    "wired-compiled": {
        "wired_gb": 42,
        "compile_blocks": True,
        "eval_every": 0,
        "keep_refiners": True,
    },
    "wired-fast": {
        "wired_gb": 42,
        "compile_blocks": True,
        "eval_every": 0,
        "keep_refiners": True,
        "vae_batch": 8,
    },
    "packed12": {
        "wired_gb": 34,
        "compile_blocks": True,
        "eval_every": 0,
        "keep_refiners": True,
        "weight_storage": "packed12",
    },
    "packed12-fast": {
        "wired_gb": 34,
        "compile_blocks": True,
        "eval_every": 0,
        "keep_refiners": True,
        "weight_storage": "packed12",
        "vae_batch": 8,
    },
    "five-pass": {
        "wired_gb": 42,
        "compile_blocks": True,
        "eval_every": 0,
        "keep_refiners": True,
        "vae_batch": 8,
    },
    "fp16": {
        "allow_nonfinite": True,
        "wired_gb": 42,
        "compile_blocks": True,
        "eval_every": 0,
        "keep_refiners": True,
        "vae_batch": 8,
    },
}


def main(a):
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S"
    )
    root = Path(a.output)
    root.mkdir(parents=True, exist_ok=True)
    cases = [
        c
        for c in json.loads(Path("experiments/cases.json").read_text())
        if c["id"] in a.cases
    ]
    for name in a.profiles:
        original_load = Linear.load.__func__
        if name == "fp16":

            def half_load(cls, store, prefix, original_load=original_load, **kwargs):
                import mlx.core as mx

                layer = original_load(cls, store, prefix, **kwargs)
                if prefix.startswith("blocks."):
                    layer = HalfProjection(
                        layer.weight.astype(mx.float16),
                        None if layer.bias is None else layer.bias.astype(mx.float16),
                    )
                return layer

            Linear.load = classmethod(half_load)
        pipeline = Pipeline(a.model, a.text_encoder, **PROFILES[name])
        try:
            for case in cases:
                for trial in range(a.repeat):
                    _, row = pipeline.generate(
                        case["prompt"],
                        width=case["width"],
                        height=case["height"],
                        seed=case["seed"],
                        images=case.get("images", ()),
                        save_trace=True,
                        steps=[0, 2, 3, 4, 5] if name == "five-pass" else None,
                        output=root / case["id"] / f"{name}-{trial}.png",
                    )
                    row.update(profile=name, case=case["id"], trial=trial)
                    if name == "fp16":
                        row["experimental_precision"] = (
                            "Main projection GEMMs FP16, native BF16 output rounding; norms, attention and other components retain the original policy"
                        )
                        Path(row["output"]).with_suffix(".json").write_text(
                            json.dumps(row, indent=2) + "\n"
                        )
                    with (root / "trials.jsonl").open("a") as f:
                        f.write(json.dumps(row) + "\n")
        finally:
            pipeline.clear()
            del pipeline
            gc.collect()
            Linear.load = classmethod(original_load)


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model", required=True)
    p.add_argument("--text-encoder", required=True)
    p.add_argument(
        "--profiles", nargs="+", choices=PROFILES, default=["compiled-batched"]
    )
    p.add_argument("--cases", nargs="+", default=["product"])
    p.add_argument("--repeat", type=int, default=3)
    p.add_argument("--output", default="outputs/optimization-v2")
    main(p.parse_args())