"""Train the mushroom body on all data and save it as JSON for the M5Stack.

    python experiments/export_for_m5stack.py data/demo_banana --intensity-pns 8 --out firmware/enose_model.json
    python experiments/export_for_m5stack.py data/demo_calibration --connectome data/connectome/kc_wiring.npy

Ripeness stages (E7) need --intensity-pns 8, the model that E7 evaluates; gas identity uses 0.

Upload enose_model.json together with firmware/mb_infer.py and firmware/enose_classifier.py
to the M5Stack. Only the non-zero PN->KC connections are stored (and no weights at all
for random 0/1 wiring). The recording schedule (heater steps, BME_PERIOD_S, purge and
sample seconds) is saved in the file and used by enose_classifier.py, so the device
measures exactly like the data the model was trained on.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from enose.common import load_table  # noqa: E402
from enose.mushroom import MushroomBodyClassifier  # noqa: E402


SCHEDULE: dict = {"heater_steps_c": [200, 240, 280, 320, 360], "heater_ms": 150, "bme_period_s": 3,
                  "purge_s": 120, "sample_s": 240}


def export(model: MushroomBodyClassifier, path: Path, schedule: dict | None = None) -> None:
    if model.kc_mode != "topk":
        raise ValueError("the M5Stack only implements kc_mode='topk' (APL); threshold mode is a laptop-only ablation")
    inputs = [np.flatnonzero(row).tolist() for row in model.wiring_]
    binary = bool(np.all((model.wiring_ == 0) | (model.wiring_ == 1)))
    payload = {
        "classes": [str(label) for label in model.classes_],
        "scale": model.scale_.tolist(),  # full precision: tiny gains must not round to 0
        "normalize": bool(model.normalize),
        "active": max(1, int(round(model.wiring_.shape[0] * model.active_fraction))),
        "kc_inputs": inputs,
        "kc_weights": None if binary else [row[index].tolist() for row, index in zip(model.wiring_, inputs)],
        "weights": model.weights_.tolist(),
        "temperature": model.temperature,
        "schedule": schedule or SCHEDULE,
        "intensity_centres": model.intensity_centres_.tolist() if model.intensity_pns else [],
        "intensity_width": float(model.intensity_width_) if model.intensity_pns else 1.0,
        "intensity_gain": float(model.intensity_gain),
        "adapt_mean": model.adapt_mean_.tolist() if model.adaptation else [],
        "adapt_std": model.adapt_std_.tolist() if model.adaptation else [],
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, separators=(",", ":")))
    print(f"已寫入 {path}（{path.stat().st_size / 1024:.0f} KB）")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("folder", type=Path)
    parser.add_argument("--out", type=Path, default=Path("firmware/enose_model.json"))
    parser.add_argument("--connectome", type=Path)
    parser.add_argument("--intensity-pns", type=int, default=0,
                        help="8 for ripeness stages (pattern + intensity), 0 for gas identity")
    parser.add_argument("--bme-period", type=int, default=SCHEDULE["bme_period_s"],
                        help="must equal BME_PERIOD_S in enose_logger.py when the training data was recorded")
    parser.add_argument("--purge", type=int, default=SCHEDULE["purge_s"])
    parser.add_argument("--sample", type=int, default=SCHEDULE["sample_s"])
    arguments = parser.parse_args()
    _, features, labels, _ = load_table(arguments.folder / "features.csv")
    connectivity = np.load(arguments.connectome) if arguments.connectome else None
    model = MushroomBodyClassifier(connectivity=connectivity, intensity_pns=arguments.intensity_pns).fit(features, labels)
    schedule = dict(SCHEDULE, bme_period_s=arguments.bme_period, purge_s=arguments.purge, sample_s=arguments.sample)
    export(model, arguments.out, schedule)


if __name__ == "__main__":
    main()