#!/usr/bin/env python3
"""Export pinned original Parakeet TDT weights; save references and provenance.

Writes a fresh portable bundle; no publication and no claimed parity until the
native reference suite has run. Requires the existing apple/.venv toolchain.

`--encoder-weights pal8` (or pal6) palettizes the encoder's weights with 8-bit (6-bit)
per-tensor k-means, spec 159's technique, keeping its 16-bit inputs and compute.
`--small-weights fp16` stores the predictor's and joint's weights in 16 bits behind
their 32-bit interface. Either makes a smaller bundle with its own adapter ID
(spec 172); the defaults reproduce the original FP16 bundle.
"""
import argparse
import hashlib
import importlib.metadata
import json
from pathlib import Path
import shutil
import time
import traceback
import numpy as np
import torch
from torch import nn
from parakeet_adapters import build_components, load_source


class StoredHalf(nn.Module):
    """16-bit weights behind a 32-bit interface: inputs cast down, outputs cast up."""
    def __init__(self, inner):
        super().__init__()
        self.inner = inner.half()

    def forward(self, *args):
        out = self.inner(*(x.half() if x.is_floating_point() else x for x in args))
        return tuple(x.float() for x in out) if isinstance(out, tuple) else out.float()


class HalfOutput(nn.Module):
    """The palettized encoder's output as 16-bit, as the FP16 encoder's was: the runner
    reads `encoded` as Float16 (Parakeet.swift's contract)."""
    def __init__(self, inner):
        super().__init__()
        self.inner = inner

    def forward(self, *args):
        return self.inner(*args).half()


def palettize(module, bits, example):
    from coreai_opt.palettization import KMeansPalettizer, KMeansPalettizerConfig
    preset = {8: KMeansPalettizerConfig.presets.w8, 6: KMeansPalettizerConfig.presets.w6}[bits]()
    palettizer = KMeansPalettizer(module, preset)
    palettizer.prepare(example, num_workers=8)
    return palettizer.finalize().eval()


def sha(path):
    with Path(path).open("rb") as f:
        return hashlib.file_digest(f, "sha256").hexdigest()


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--source", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--resume", action="store_true", help="Keep already-exported components after an interrupted export")
    p.add_argument("--encoder-weights", choices=["fp16", "pal8", "pal6"], default="fp16")
    p.add_argument("--small-weights", choices=["fp32", "fp16"], default="fp32", help="Predictor and joint weights")
    a = p.parse_args()
    slim = a.encoder_weights != "fp16" or a.small_weights != "fp32"
    adapter_id = "parakeet-tdt-v3-coreai-batch15-v1" if not slim else f"parakeet-tdt-v3-coreai-{a.encoder_weights}-{a.small_weights}-v1"
    if a.output.exists() and not a.resume:
        p.error("Use a fresh output directory, or explicitly --resume the same source export")
    a.output.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(4)
    torch.manual_seed(1729)
    provenance = json.loads((a.source / "source-provenance.json").read_text())
    provenance["weights_sha256"] = sha(a.source / "model.safetensors")
    report_path = a.output / "conversion.json"
    previous = json.loads(report_path.read_text()) if report_path.exists() else None
    if previous and previous["source"] != provenance:
        raise ValueError("Cannot resume an export from different source weights")
    report = previous or {"source": provenance, "status": "started", "components": {}, "weights": {"encoder": a.encoder_weights, "predictor_joint": a.small_weights},
        "versions": {n: importlib.metadata.version(n) for n in ["torch", "transformers", "coreai-torch", "coreai-core"]}}
    def save():
        report_path.write_text(json.dumps(report, indent=2, allow_nan=False)+"\n")
    save()
    try:
        model = load_source(a.source)
        from coreai_torch import TorchConverter, get_decomp_table
        for name, (graph, inputs, names, outputs, precision) in build_components(model).items():
            if report["components"].get(name, {}).get("status") == "exported":
                continue
            start = time.perf_counter()
            entry = {"status": "started", "precision": precision}
            report["components"][name] = entry
            save()
            # Encoder/projector precision does not modify predictor or joint modules.
            graph = graph.eval()
            if name == "encoder" and a.encoder_weights != "fp16":
                graph = HalfOutput(palettize(graph, int(a.encoder_weights[3:]), inputs))
                precision = a.encoder_weights + "-fp16"
                entry["precision"] = precision
            if name in ("predictor", "joint") and a.small_weights == "fp16":
                graph = StoredHalf(graph).eval()
                precision = "fp16-weights"
                entry["precision"] = precision
            if precision.endswith("fp16") and name == "encoder":
                graph.half()
                inputs = tuple(x.half() if x.is_floating_point() else x for x in inputs)
            with torch.no_grad():
                result = graph(*inputs)
                result = result if isinstance(result, tuple) else (result,)
                if not all(torch.isfinite(x).all() for x in result):
                    raise ValueError("Nonfinite PyTorch reference: " + name)
                reference = a.output / "references"
                reference.mkdir(exist_ok=True)
                np.savez(reference/(name+".npz"), **{n: x.numpy() for n,x in zip(names,inputs)}, **{"expected_"+n:x.numpy() for n,x in zip(outputs,result)})
                # Portable raw fixtures consumed by SwiftParakeetComponentBench.
                for tensor_name, tensor in [*zip(names, inputs), *(("expected_"+n, x) for n,x in zip(outputs,result))]:
                    tensor.float().numpy().astype("<f4").tofile(reference/(name+"-"+tensor_name+".f32"))
                entry.update(status="reference_ready", inputs=[{"name":n,"shape":list(x.shape),"dtype":str(x.dtype).removeprefix("torch.")} for n,x in zip(names,inputs)], outputs=outputs)
                save()
                ep = torch.export.export(graph, args=inputs).run_decompositions(get_decomp_table())
                entry["status"] = "torch_exported"
                save()
                program = TorchConverter(mode=TorchConverter.Mode.RELEASE).add_exported_program(ep, input_names=names, output_names=outputs).to_coreai()
                entry["status"] = "coreai_lowered"
                save()
                program.optimize()
                program.save_asset(a.output/(name+".aimodel"))
            entry.update(status="exported", seconds=time.perf_counter()-start)
            save()
            print(json.dumps({"component":name, **entry}),flush=True)
        shutil.copy2(a.source/"tokenizer.json", a.output/"tokenizer.json")
        shutil.copy2(a.source/"README.md", a.output/"SOURCE_MODEL_CARD.md")
        files=[]
        for asset in sorted(a.output.glob("*.aimodel")):
            for f in sorted(asset.rglob("*")):
                if f.is_file(): files.append({"path":str(f.relative_to(a.output)),"bytes":f.stat().st_size,"sha256":sha(f)})
        for f in [a.output/"tokenizer.json",a.output/"SOURCE_MODEL_CARD.md"]:
            files.append({"path":f.name,"bytes":f.stat().st_size,"sha256":sha(f)})
        manifest={"adapter_id":adapter_id,"source":provenance,"license":"CC-BY-4.0","sample_rate":16000,"max_window_samples":240000,"max_utterance_samples":560000,"blank_id":8192,"durations":[0,1,2,3,4],"encoder_frames":188,"encoder_width":640,"frontend_frames":1500,"components":{name:{"asset":name+".aimodel","inputs":v["inputs"],"outputs":v["outputs"],"precision":v["precision"]} for name,v in report["components"].items()},"files":files,"validation":"pending_native_parity"}
        (a.output/"manifest.json").write_text(json.dumps(manifest,indent=2)+"\n")
        report["status"]="exported"
        save()
    except BaseException as error:
        report.update(status="failed",error=str(error),traceback=traceback.format_exc())
        save()
        raise


if __name__ == "__main__":
    main()