mnn_export.py

"""
mnn_export.py
=============
"""

from __future__ import annotations

import argparse
import hashlib
import logging
import os
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import zone_observation as _zo

assert _zo.SCHEMA_VERSION == 3, (
    f"mnn_export: zone_observation schema mismatch "
    f"(expected 3, got {_zo.SCHEMA_VERSION})"
)

from zone_observation import ForecastConfig

try:
    from weather_forecast_env import make_weather_env
    from gru_weather_policy import GRUWeatherFeaturesExtractor, create_gru_weather_policy_kwargs
    _ML_AVAILABLE = True
except ImportError:
    _ML_AVAILABLE = False
    GRUWeatherFeaturesExtractor      = None  # type: ignore[assignment,misc]
    create_gru_weather_policy_kwargs = None  # type: ignore[assignment]
    make_weather_env                 = None  # type: ignore[assignment]

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Hardware capability reporter
# ---------------------------------------------------------------------------

def report_edge_capability() -> str:
    import subprocess

    try:
        out = subprocess.run(
            ["vulkaninfo", "--summary"],
            capture_output=True, text=True, timeout=5,
        ).stdout
        if "Vulkan" in out:
            logger.info("Edge capability: Vulkan detected -> MNN Vulkan backend")
            return "VULKAN"
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    try:
        out = subprocess.run(
            ["clinfo"], capture_output=True, text=True, timeout=5,
        ).stdout
        if "Mali" in out:
            if "OpenCL 2" in out or "OpenCL 3" in out:
                logger.info("Edge capability: Mali + OpenCL 2/3 -> MNN OpenCL backend")
                return "OPENCL"
            if "Mali-450" in out or "Utgard" in out:
                logger.warning(
                    "Edge capability: Mali-450 (Utgard) detected. "
                    "No OpenCL / Vulkan support. Forcing CPU fallback. "
                    "The quantized .mnn file will still run — just slower."
                )
                return "CPU_ONLY"
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    logger.info("Edge capability: GPU info unavailable -> defaulting to CPU_ONLY")
    return "CPU_ONLY"


# ---------------------------------------------------------------------------
# StatelessInferenceWrapper — explicit hidden state I/O for GRU tracing
# ---------------------------------------------------------------------------

class StatelessInferenceWrapper:

    def __init__(
        self,
        features_extractor,
        mlp_extractor,
        action_net,
        hidden_size: int,
        obs_keys:    List[str],
    ):
        self.features_extractor = features_extractor
        self.mlp_extractor      = mlp_extractor
        self.action_net         = action_net
        self.hidden_size        = hidden_size
        self.obs_keys           = obs_keys

    def forward(
        self,
        obs_tensors: List,   # one tensor per obs_key, in canonical order
        hidden_in,           # [1, 1, hidden_size]
    ) -> Tuple:
        # Rebuild obs dict from positional tensors (tracing-safe)
        obs = {k: obs_tensors[i] for i, k in enumerate(self.obs_keys)}

        # Inject external hidden state into the features extractor
        self.features_extractor.set_hidden(hidden_in)

        # Extract features (MLP/conv + GRU step)
        features = self.features_extractor(obs)

        # Actor path only — discard value/critic at export time
        latent_pi = self.mlp_extractor.forward_actor(features)

        # Action logits [1, n_actions]
        action_logits = self.action_net(latent_pi)

        # Return updated hidden state for the edge runtime to store
        hidden_out = self.features_extractor.get_hidden()

        return action_logits, hidden_out


# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------

def _validate_output_path(path_str: str) -> Path:
    p = Path(path_str).resolve()
    if p.suffix != ".mnn":
        raise ValueError(
            f"Output path must end with .mnn, got: {path_str!r}"
        )
    p.parent.mkdir(parents=True, exist_ok=True)
    return p


def _sha256_file(path: Path, chunk_size: int = 1 << 20) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(chunk_size):
            h.update(chunk)
    return h.hexdigest()


# ---------------------------------------------------------------------------
# Calibration data collection
# ---------------------------------------------------------------------------

def _collect_calibration_obs(
    env,
    n_episodes: int = 100,
    obs_keys:   Optional[List[str]] = None,
) -> List[Dict]:
    import numpy as np

    logger.info("Collecting calibration data (%d episodes)...", n_episodes)
    samples = []
    for ep in range(n_episodes):
        obs, _ = env.reset()
        done = False
        steps = 0
        while not done and steps < 50:
            sample = {k: v for k, v in obs.items() if k != "action_mask"}
            samples.append(sample)
            action = env.action_space.sample()
            obs, _, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            steps += 1
        if (ep + 1) % 20 == 0:
            logger.info("  Calibration: %d/%d episodes", ep + 1, n_episodes)

    logger.info("Collected %d calibration samples", len(samples))
    return samples


# ---------------------------------------------------------------------------
# ONNX export
# ---------------------------------------------------------------------------

def _export_onnx(
    wrapper:      StatelessInferenceWrapper,
    dummy_obs:    Dict,
    hidden_size:  int,
    onnx_path:    Path,
    export_keys:  List[str],
) -> None:
    import torch

    obs_tensors = [dummy_obs[k].float() for k in export_keys]
    dummy_hidden = torch.zeros(1, 1, hidden_size)

    input_names  = export_keys + ["hidden_in"]
    output_names = ["action_logits", "hidden_out"]

    dynamic_axes: Dict[str, Dict[int, str]] = {k: {0: "batch"} for k in export_keys}
    dynamic_axes["hidden_in"]      = {1: "batch"}
    dynamic_axes["action_logits"]  = {0: "batch"}
    dynamic_axes["hidden_out"]     = {1: "batch"}

    logger.info("Exporting to ONNX (opset 17): %s", onnx_path)
    logger.info("  Observation inputs: %s", export_keys)
    logger.info("  Outputs: %s", output_names)

    with torch.no_grad():
        torch.onnx.export(
            wrapper,
            args=(obs_tensors, dummy_hidden),
            f=str(onnx_path),
            opset_version=17,
            input_names=input_names,
            output_names=output_names,
            dynamic_axes=dynamic_axes,
        )

    logger.info(
        "ONNX saved: %s  (%.1f MB)",
        onnx_path, onnx_path.stat().st_size / 1e6,
    )


# ---------------------------------------------------------------------------
# MNN conversion
# ---------------------------------------------------------------------------

def _convert_to_mnn(
    onnx_path:             Path,
    mnn_path:              Path,
    quantize:              str,
    calibration_samples:   Optional[List[Dict]] = None,
) -> None:
    logger.info(
        "Converting to MNN  quantize=%s  target=%s", quantize, mnn_path
    )
    converted = False

    try:
        from MNN.tools import mnnconvert as _mnnconvert
        args = {
            "modelFile": str(onnx_path),
            "MNNModel":  str(mnn_path),
            "framework": "ONNX",
            "bizCode":   "weather_rl_v1",
        }
        if quantize == "int8":
            args["weightQuantBits"] = 8
        elif quantize == "fp16":
            args["fp16"] = True
        _mnnconvert.convert(args)
        converted = True
        logger.info("MNN conversion via Python API: OK")
    except Exception as api_err:
        logger.warning("MNN Python API failed (%s) — trying CLI fallback", api_err)

    if not converted:
        import subprocess, shutil
        cli = shutil.which("mnnconvert")
        if cli is None:
            raise RuntimeError(
                "mnnconvert not found on PATH and MNN Python API failed.\n"
                "Build from: https://github.com/GeniusVentures/MNN\n"
                "Or install: pip install MNN"
            )
        cmd = [cli, "-f", "ONNX",
               "--modelFile", str(onnx_path),
               "--MNNModel",  str(mnn_path)]
        if quantize == "int8":
            cmd += ["--weightQuantBits", "8"]
        elif quantize == "fp16":
            cmd += ["--fp16"]
        logger.info("MNN CLI: %s", " ".join(cmd))
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            raise RuntimeError(
                f"mnnconvert CLI failed (rc={result.returncode}):\n"
                f"stdout: {result.stdout}\nstderr: {result.stderr}"
            )

    if not mnn_path.exists():
        raise RuntimeError(
            f"MNN conversion reported success but {mnn_path} was not created."
        )


# ---------------------------------------------------------------------------
# Main export function
# ---------------------------------------------------------------------------

def export_to_mnn(
    checkpoint_path:       str,
    output_mnn:            str  = "weather_rl_model.mnn",
    quantize:              str  = "int8",
    calibration_episodes:  int  = 0,
    hidden_size:           int  = 64,
    n_zones:               int  = 4,
    keep_onnx:             bool = False,
) -> Path:
    import torch
    from sb3_contrib import MaskablePPO

    ckpt = Path(checkpoint_path)
    if not ckpt.exists():
        raise FileNotFoundError(f"Checkpoint not found: {ckpt}")
    if quantize not in ("int8", "fp16", "none"):
        raise ValueError(f"quantize must be 'int8', 'fp16', or 'none', got {quantize!r}")

    mnn_path  = _validate_output_path(output_mnn)
    onnx_path = mnn_path.with_suffix(".onnx")

    backend = report_edge_capability()

    # --- Load MaskablePPO checkpoint ---
    logger.info("Loading checkpoint: %s", ckpt)
    if n_zones < 1:
        raise ValueError(f"n_zones must be >= 1, got {n_zones}")
    config = ForecastConfig(n_zones=n_zones, horizon_days=30)
    env    = make_weather_env(config)
    logger.info("Export env: n_zones=%d  horizon_days=30", n_zones)

    custom_objects = {}
    if GRUWeatherFeaturesExtractor is not None:
        custom_objects = {
            "features_extractor_class": GRUWeatherFeaturesExtractor,
        }

    model = MaskablePPO.load(
        str(ckpt),
        env=env,
        device="cpu",
        custom_objects=custom_objects if custom_objects else None,
    )
    model.policy.eval()

    policy = model.policy
    assert hasattr(policy, "mlp_extractor") and hasattr(policy, "action_net"), (
        f"Loaded policy is missing expected actor components. "
        f"Got attributes: {[a for a in dir(policy) if not a.startswith('_')]}"
    )
    assert hasattr(policy, "features_extractor"), (
        "Loaded policy is missing features_extractor."
    )
    logger.info("Checkpoint loaded: %s", ckpt.name)

    _ckpt_hidden = getattr(policy.features_extractor, "hidden_size", None)
    if _ckpt_hidden is not None and _ckpt_hidden != hidden_size:
        logger.warning(
            "Overriding --hidden-size=%d with checkpoint's hidden_size=%d",
            hidden_size, _ckpt_hidden,
        )
        hidden_size = _ckpt_hidden

    features_extractor = policy.features_extractor
    mlp_extractor      = policy.mlp_extractor
    action_net         = policy.action_net

    obs_sample, _ = env.reset()
    obs_keys_all  = sorted(obs_sample.keys())
    export_keys   = [k for k in obs_keys_all if k != "action_mask"]

    wrapper = StatelessInferenceWrapper(
        features_extractor=features_extractor,
        mlp_extractor=mlp_extractor,
        action_net=action_net,
        hidden_size=hidden_size,
        obs_keys=export_keys,
    )

    # --- Build dummy input ---
    dummy_obs: Dict[str, "torch.Tensor"] = {}
    for k, v in obs_sample.items():
        t = torch.from_numpy(v).unsqueeze(0)
        dummy_obs[k] = t.float() if k != "action_mask" else t

    # --- Optional calibration ---
    calibration_samples = None
    if calibration_episodes > 0 and quantize == "int8":
        calibration_samples = _collect_calibration_obs(
            env, n_episodes=calibration_episodes
        )

    # --- ONNX export ---
    _export_onnx(wrapper, dummy_obs, hidden_size, onnx_path, export_keys)

    # --- MNN conversion ---
    try:
        _convert_to_mnn(onnx_path, mnn_path, quantize, calibration_samples)
    finally:
        if onnx_path.exists() and not keep_onnx:
            onnx_path.unlink()
            logger.info("Removed intermediate ONNX: %s", onnx_path.name)

    if not mnn_path.exists():
        raise RuntimeError(
            f"Export appeared to succeed but {mnn_path} was not created."
        )

    size_mb = mnn_path.stat().st_size / (1024 * 1024)
    sha     = _sha256_file(mnn_path)

    logger.info(
        "MNN export complete: %s  (%.1f MB)  SHA256: %s", mnn_path, size_mb, sha
    )
    logger.info("Edge backend detected: %s", backend)

    print(f"""
╔══════════════════════════════════════════════════════════╗
║  Weather RL Model — Edge Deployment Manifest             ║
╠══════════════════════════════════════════════════════════╣
║  Model:      {mnn_path.name:<44} ║
║  Size:       {f'{size_mb:.1f} MB':<44} ║
║  Quantize:   {quantize:<44} ║
║  Backend:    {backend:<44} ║
║  SHA256:     {sha[:44]}  ║
║              {sha[44:]}  ║
╠══════════════════════════════════════════════════════════╣
║  POST-PROCESSING (apply in edge runtime):                ║
║    logits = model.run(obs_without_mask, hidden_in)       ║
║    logits[action_mask == 0] = -1e9                       ║
║    action = argmax(logits)                               ║
║    Store hidden_out; pass as hidden_in next step         ║
╚══════════════════════════════════════════════════════════╝
""")

    return mnn_path


# ---------------------------------------------------------------------------
# Edge inference protocol (copy into edge runtime documentation)
# ---------------------------------------------------------------------------
EDGE_INFERENCE_NOTE = """
Edge Runtime Inference Protocol
================================
The exported .mnn model is a stateless actor network. The edge runtime
must manage two pieces of state externally:

1. GRU hidden state (temporal belief):
   - Initialise: hidden = zeros([1, 1, {hidden_size}])
   - Each step:  action_logits, hidden = model.run(obs_inputs, hidden)
   - Reset:      hidden = zeros([1, 1, {hidden_size}]) at episode start

2. Action mask (zone validity):
   - The model outputs raw action logits [1, n_zones + 1]
   - Apply mask BEFORE argmax:
       action_logits[action_mask == 0] = -1e9
       action = argmax(action_logits)
   - The terminate action (index n_zones) is ALWAYS valid; never mask it.

Input tensor order (must match ONNX input_names exactly):
   {obs_keys_without_mask}  (float32)
   hidden_in                (float32, shape [1, 1, H])

Output tensors:
   action_logits  float32  [1, n_zones + 1]  raw scores; apply mask + argmax
   hidden_out     float32  [1, 1, H]         store and feed back next step

Note (schema v3): the observation inputs now include basin_context
(float32, shape [1, 8] -- [enso_oni, iod_dmi, itcz_latitude_deg,
mslp_regional_hpa, solar_wind_speed_kms, kp_index, goes_xray_log10,
helio_regime_ord] in that field order, matching
weather_forecast_env.basin_context_vector() -- NOT a BasinContext method;
BasinContext itself only exposes to_dict()/from_dict()).
It sorts FIRST in the alphabetical input order above. edge_wrapper.cpp
already matches (BASIN_FLAT=8, verified against the same 8-value neutral
default below); any other edge runtime built against the pre-v3 (4-input,
climate-only) interface MUST add the four helio inputs -- the MNN session
will fail or produce garbage logits if basin_context is left unbound at
the old width. When no basin data is available at the edge, feed the
neutral default (0.0, 0.0, 0.0, 1013.25, 400.0, 2.0, -7.0, 0.0) -- the
same 8-value default weather_forecast_env._BASIN_CONTEXT_NEUTRAL and
edge_wrapper.cpp's BASIN_NEUTRAL both use.
"""


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Export MaskablePPO weather policy to quantized .mnn",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("--checkpoint",  required=True,
                   help="Path to trained .zip checkpoint")
    p.add_argument("--output",      default="weather_rl_model.mnn",
                   help="Output .mnn path")
    p.add_argument("--quantize",    default="int8",
                   choices=["int8", "fp16", "none"])
    p.add_argument("--calibration-episodes", type=int, default=0,
                   help="Episodes for PTQ calibration (0 = weight-only)")
    p.add_argument("--hidden-size", type=int, default=64)
    p.add_argument(
        "--n-zones", type=int, default=4,
        help=(
            "Zones the checkpoint was trained with. "
            "normal=2, monsoon/drought=3, heatwave/humidity=4 (default 4). "
            "Must match the curriculum phase or the ONNX trace will have wrong "
            "input shapes and be incompatible with edge_wrapper.cpp."
        ),
    )
    p.add_argument("--keep-onnx",   action="store_true")
    p.add_argument("--capability",  action="store_true",
                   help="Report edge hardware capability and exit")
    return p.parse_args()


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(message)s",
    )
    args = _parse_args()

    if args.capability:
        print(f"Edge backend: {report_edge_capability()}")
        sys.exit(0)

    try:
        export_to_mnn(
            checkpoint_path=args.checkpoint,
            output_mnn=args.output,
            quantize=args.quantize,
            calibration_episodes=args.calibration_episodes,
            hidden_size=args.hidden_size,
            n_zones=args.n_zones,
            keep_onnx=args.keep_onnx,
        )
        sys.exit(0)
    except FileNotFoundError as e:
        logger.error("Checkpoint not found: %s", e)
        sys.exit(2)
    except ValueError as e:
        logger.error("Invalid argument: %s", e)
        sys.exit(2)
    except RuntimeError as e:
        logger.error("Export failed: %s", e)
        sys.exit(1)
    except Exception as e:
        logger.exception("Unexpected error: %s", e)
        sys.exit(1)


# ---------------------------------------------------------------------------
# Self-test (no checkpoint or MNN required)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    if "--checkpoint" in sys.argv:
        main()
    else:
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s | %(levelname)s | %(message)s",
        )
        print("=== mnn_export.py self-test (no checkpoint/MNN required) ===\n")
        failures = []

        def _assert(cond: bool, msg: str) -> None:
            if not cond:
                failures.append(msg)
                print(f"  FAIL: {msg}")

        # 1. _validate_output_path rejects non-.mnn extensions
        try:
            _validate_output_path("/tmp/model.pkl")
            _assert(False, "Should have rejected .pkl extension")
        except ValueError:
            pass
        try:
            p = _validate_output_path("/tmp/test_export.mnn")
            _assert(p.suffix == ".mnn", "Resolved path should end in .mnn")
        except Exception as e:
            _assert(False, f"Valid .mnn path rejected: {e}")
        print("  _validate_output_path OK")

        # 2. StatelessInferenceWrapper can be constructed with mock components
        try:
            import torch
            import torch.nn as nn

            class _FakeExtractor:
                def __call__(self, obs): return torch.zeros(1, 256)
                def set_hidden(self, h): self._h = h
                def get_hidden(self): return getattr(self, '_h', torch.zeros(1,1,64))

            class _FakeMLPExtractor(nn.Module):
                def forward_actor(self, x): return x[:, :128]

            wrapper = StatelessInferenceWrapper(
                features_extractor=_FakeExtractor(),
                mlp_extractor=_FakeMLPExtractor(),
                action_net=nn.Linear(128, 5),
                hidden_size=64,
                obs_keys=["basin_context", "forecast_precip",
                          "forecast_uncertainty", "prior_belief",
                          "zone_belief"],
            )
            _assert(wrapper.hidden_size == 64, "Wrong hidden_size on wrapper")
            _assert(len(wrapper.obs_keys) == 5, "Wrong obs_keys count")
            print("  StatelessInferenceWrapper construction OK")
        except ImportError as e:
            print(f"  StatelessInferenceWrapper: torch not installed, skipped ({e})")

        # 3. _sha256_file is deterministic
        import tempfile
        with tempfile.NamedTemporaryFile(delete=False, suffix=".bin") as f:
            f.write(b"weather_rl_test" * 1000)
            tmp = Path(f.name)
        sha1 = _sha256_file(tmp)
        sha2 = _sha256_file(tmp)
        _assert(sha1 == sha2, "SHA256 not deterministic")
        _assert(len(sha1) == 64, f"SHA256 wrong length: {len(sha1)}")
        tmp.unlink()
        print(f"  _sha256_file OK  sha={sha1[:16]}...")

        # 4. report_edge_capability returns a known string
        cap = report_edge_capability()
        _assert(cap in ("VULKAN", "OPENCL", "CPU_ONLY"),
                f"Unknown capability: {cap!r}")
        print(f"  report_edge_capability OK  backend={cap}")

        # 5. SCHEMA_VERSION guard
        _assert(_zo.SCHEMA_VERSION == 3,
                f"SCHEMA_VERSION guard not working (got {_zo.SCHEMA_VERSION})")
        print("  SCHEMA_VERSION guard OK")

        # 6. EDGE_INFERENCE_NOTE is complete
        _assert(len(EDGE_INFERENCE_NOTE) > 100,    "EDGE_INFERENCE_NOTE too short")
        _assert("hidden_out" in EDGE_INFERENCE_NOTE, "missing hidden_out")
        _assert("argmax" in EDGE_INFERENCE_NOTE,     "missing argmax")
        _assert("action_logits" in EDGE_INFERENCE_NOTE, "missing action_logits")
        print("  EDGE_INFERENCE_NOTE present and complete")

        # 7. export_keys excludes action_mask (schema v3 key set incl. basin_context)
        sample_keys = ["action_mask", "basin_context", "forecast_precip",
                       "forecast_uncertainty", "prior_belief", "zone_belief"]
        export = [k for k in sorted(sample_keys) if k != "action_mask"]
        _assert("action_mask" not in export, "action_mask leaked into export_keys")
        _assert(len(export) == 5, f"Expected 5 export keys, got {len(export)}")
        _assert(export[0] == "basin_context",
                "basin_context should sort first (alphabetical)")
        print("  export_keys exclusion of action_mask OK")

        print()
        if failures:
            print(f"FAILED  {len(failures)} test(s):")
            for f in failures:
                print(f"  - {f}")
            sys.exit(1)
        else:
            print("All mnn_export self-tests passed.")
            print()
            print("To run the real export:")
            print("  python mnn_export.py --checkpoint final_normal.zip")
            print("  python mnn_export.py --capability")