"""Hand-written CiteNet -> Core ML (.mlpackage, fp16), plus parity and placement checks."""
import json, os, shutil, sys, time
import numpy as np, torch, coremltools as ct
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from citenet import net as cn, text, spec
from lepc import load

OUT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "build")
# Fixed shapes: with EnumeratedShapes Core ML ran every shape on the CPU. With a
# fixed batch the planner puts all ops on the ANE from B=16 up; B=1 stays on CPU.
# b2048 / b8192 are for the GPU, which wants bigger batches.
BATCHES = [1, 32, 512, 2048, 8192]

model, meta = cn.make()
tot, nz = cn.param_counts(model)
paths = {}
for B in BATCHES:
    traced = torch.jit.trace(model, torch.zeros(B, 1, 1, spec.SEQ))
    ml = ct.convert(
        traced,
        inputs=[ct.TensorType(name="codes", shape=(B, 1, 1, spec.SEQ), dtype=np.float16)],
        outputs=[ct.TensorType(name=n, dtype=np.float16) for n in ("p_fake", "evidence", "concepts")],
        compute_precision=ct.precision.FLOAT16,
        minimum_deployment_target=ct.target.macOS15,
    )
    ml.short_description = "CiteNet: phantom legal-citation detector with hand-written weights (no training)"
    ml.user_defined_metadata.update({"params": str(tot), "nonzero_params": str(nz), "trained": "false",
                                     "batch": str(B), "evidence": json.dumps(meta["evidence"])})
    paths[B] = os.path.join(OUT, f"CiteNet_b{B}.mlpackage")
    ml.save(paths[B])
    print("saved", paths[B])
# lean builds for bulk scanning: no concepts output (only p_fake + evidence leave the accelerator)
class Lean(torch.nn.Module):
    def __init__(self, m):
        super().__init__(); self.m = m
    def forward(self, x):
        p, e, _ = self.m(x)
        return p, e
for B in (512, 2048, 8192):
    traced = torch.jit.trace(Lean(model).eval(), torch.zeros(B, 1, 1, spec.SEQ))
    ml = ct.convert(traced, inputs=[ct.TensorType(name="codes", shape=(B, 1, 1, spec.SEQ), dtype=np.float16)],
                    outputs=[ct.TensorType(name=n, dtype=np.float16) for n in ("p_fake", "evidence")],
                    compute_precision=ct.precision.FLOAT16, minimum_deployment_target=ct.target.macOS15)
    ml.short_description = "CiteNet lean (p_fake + evidence only), hand-written weights"
    lp = os.path.join(OUT, f"CiteNet_lean_b{B}.mlpackage"); ml.save(lp)
    dst = os.path.join(OUT, f"CiteNet_lean_b{B}.mlmodelc"); shutil.rmtree(dst, ignore_errors=True)
    ct.utils.compile_model(lp, dst)
    print("saved", lp)
json.dump({k: meta[k] for k in ("features", "concepts", "rules", "evidence", "groups", "anchored", "floating")},
          open(os.path.join(OUT, "meta.json"), "w"))
np.savez_compressed(os.path.join(OUT, "weights.npz"), **{k: v for k, v in cn.build()[0].items()})

# ---- parity: torch fp32 vs Core ML fp16 on every train+eval citation
rows = load("aux_train") + load("eval")
X = text.encode([r["cite"] for r in rows])
with torch.no_grad():
    pt = model(torch.from_numpy(X))[0][:, 0].numpy()
for B, cu in ((512, "CPU_AND_NE"), (2048, "CPU_AND_GPU")):
    m = ct.models.MLModel(paths[B], compute_units=getattr(ct.ComputeUnit, cu))
    outs = []
    for i in range(0, len(X), B):
        xb = X[i:i + B]; n = len(xb)
        pad = np.zeros((B, 1, 1, spec.SEQ), np.float16); pad[:n] = xb
        outs.append(m.predict({"codes": pad})["p_fake"][:n, 0])
    pc = np.concatenate(outs).astype(np.float32)
    flip = np.where((pc >= 0.5) != (pt >= 0.5))[0]
    print(f"{cu} b{B} vs torch fp32: {len(X)-len(flip)}/{len(X)} labels agree; flips: {[rows[i]['cite'] for i in flip]}")

# ---- placement
from coremltools.models.compute_plan import MLComputePlan
for B, p in paths.items():
    dst = os.path.join(OUT, f"CiteNet_b{B}.mlmodelc")
    shutil.rmtree(dst, ignore_errors=True)        # compile_model will not overwrite
    c = ct.utils.compile_model(p, dst)
    for cu in ("CPU_AND_NE", "CPU_AND_GPU"):
        plan = MLComputePlan.load_from_path(path=c, compute_units=getattr(ct.ComputeUnit, cu))
        cnt = {}
        for fn in plan.model_structure.program.functions.values():
            for op in fn.block.operations:
                u = plan.get_compute_device_usage_for_mlprogram_operation(op)
                if u:
                    k = type(u.preferred_compute_device).__name__.replace("ML", "").replace("ComputeDevice", "")
                    cnt[k] = cnt.get(k, 0) + 1
        print(f"b{B} {cu} placement:", cnt)