#!/usr/bin/env python3
"""
DM-JEPA 1.5 Apple Neural Engine (ANE) Exporter.
Exports the 94.24M Cartesian Block Butterfly Decision Architecture to Apple Core ML (.mlpackage).
Optimized specifically for the 16-core Apple Neural Engine on iPhone (A16, A17 Pro, A18) and iPad/Mac (M1-M4).
"""

import os
import argparse
import numpy as np
import torch
import coremltools as ct
from ane_model import DMJEPA15ANEModel


def export_coreml(
    checkpoint_path: str = "checkpoints/djepa_v2.0_autonomous_production.pt",
    output_path: str = "dm_jepa_1_5_ios_ane/DMJEPA15_ANE.mlpackage",
    seq_len: int = 64,
    num_options: int = 4,
    opt_len: int = 32
):
    print("=" * 70)
    print("🚀 Exporting DM-JEPA 1.5 to Apple Neural Engine (Core ML .mlpackage)")
    print("=" * 70)

    # 1. Initialize PyTorch model
    print("Initializing ANE-optimized PyTorch model...")
    model = DMJEPA15ANEModel().eval()

    # 2. Load weights
    print(f"Loading weights from {checkpoint_path}...")
    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    sd = ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt

    ane_sd = {}
    for k, v in sd.items():
        k_ane = k
        k_ane = k_ane.replace("criteria_encoder.proj.0.", "criteria_encoder.proj_0.")
        k_ane = k_ane.replace("criteria_encoder.proj.2.", "criteria_encoder.proj_2.")
        k_ane = k_ane.replace("predictor.cell.ffn.0.", "predictor_cell.ffn_0.")
        k_ane = k_ane.replace("predictor.cell.ffn.2.", "predictor_cell.ffn_2.")
        k_ane = k_ane.replace("predictor.cell.", "predictor_cell.")
        k_ane = k_ane.replace("predictor.out_proj.0.", "predictor_out_0.")
        k_ane = k_ane.replace("predictor.out_proj.1.", "predictor_out_1.")
        k_ane = k_ane.replace("predictor.out_proj.2.", "predictor_out_2.")
        ane_sd[k_ane] = v

    missing, unexpected = model.load_state_dict(ane_sd, strict=False)
    real_missing = [m for m in missing if not m.endswith(".cos") and not m.endswith(".sin")]
    if real_missing:
        print(f"⚠️ Warning: Missing keys: {real_missing}")
    else:
        print("✅ All 181 production parameters loaded cleanly into ANE graph!")

    # 3. Dummy inputs for tracing
    state_ids = torch.randint(0, 1000, (1, seq_len), dtype=torch.int32)
    state_mask = torch.ones((1, seq_len), dtype=torch.int32)
    opt_ids = torch.randint(0, 1000, (1, num_options, opt_len), dtype=torch.int32)
    opt_mask = torch.ones((1, num_options, opt_len), dtype=torch.int32)

    print("Tracing execution graph with PyTorch JIT...")
    traced = torch.jit.trace(model, (state_ids.long(), state_mask.long(), opt_ids.long(), opt_mask.long()))

    # 4. Core ML Conversion
    inputs = [
        ct.TensorType(name="state_input_ids", shape=(1, seq_len), dtype=np.int32),
        ct.TensorType(name="state_attention_mask", shape=(1, seq_len), dtype=np.int32),
        ct.TensorType(name="option_input_ids", shape=(1, num_options, opt_len), dtype=np.int32),
        ct.TensorType(name="option_attention_mask", shape=(1, num_options, opt_len), dtype=np.int32),
    ]

    outputs = [
        ct.TensorType(name="logits"),
        ct.TensorType(name="probabilities"),
        ct.TensorType(name="metric_energy"),
    ]

    print("Converting PyTorch MIL graph to Apple Neural Engine program...")
    mlmodel = ct.convert(
        traced,
        inputs=inputs,
        outputs=outputs,
        minimum_deployment_target=ct.target.iOS16,
        compute_precision=ct.precision.FLOAT16
    )

    # 5. Metadata for Xcode
    mlmodel.author = "DangerLabs"
    mlmodel.license = "Apache 2.0"
    mlmodel.short_description = "DM-JEPA 1.5 Real-Time System 1 Decision Model for Apple Neural Engine (ANE)."
    mlmodel.version = "1.5.0"
    
    mlmodel.input_description["state_input_ids"] = "State / Instruction token IDs [1, 64]"
    mlmodel.input_description["state_attention_mask"] = "State attention mask [1, 64]"
    mlmodel.input_description["option_input_ids"] = "Candidate action token IDs [1, 4, 32]"
    mlmodel.input_description["option_attention_mask"] = "Candidate action attention mask [1, 4, 32]"

    mlmodel.output_description["logits"] = "Raw compatibility logits for candidate options [1, 4]"
    mlmodel.output_description["probabilities"] = "Softmax normalized choice probabilities [1, 4]"
    mlmodel.output_description["metric_energy"] = "Riemannian manifold metric energy [1, 4]"

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    print(f"Saving compiled Core ML package to {output_path}...")
    mlmodel.save(output_path)
    print("=" * 70)
    print("✨ Core ML package export for Apple Neural Engine complete! ✨")
    print(f"📦 Output: {output_path}")
    print("=" * 70)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", default="checkpoints/djepa_v2.0_autonomous_production.pt")
    parser.add_argument("--output", default="dm_jepa_1_5_ios_ane/DMJEPA15_ANE.mlpackage")
    args = parser.parse_args()
    export_coreml(args.checkpoint, args.output)