"""Vorhersage-Skript (Lieferumfang): laedt das gelieferte Modell in einem frischen Prozess.

  python vorhersage.py tabelle --ordner DIR --csv eingabe.csv --ausgabe vorhersagen.csv [--format torchscript|onnx]
  python vorhersage.py bilder  --ordner DIR --npy bilder.npy [--ausgabe vorhersagen.csv] [--format torchscript|onnx]

tabelle: CSV mit den Spalten alter, einkommen, nutzung_h, rauschen1, rauschen2, vertrag, region (leere Zellen erlaubt).
bilder:  .npy (n, 1, 8, 8) Werte 0..1.
Ausgabe: Zeile, Klasse, Wahrscheinlichkeit; zusaetzlich logits als .logits.npy neben der Ausgabe.
"""
from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path

import numpy as np

NUM = ["alter", "einkommen", "nutzung_h", "rauschen1", "rauschen2"]
KAT = ["vertrag", "region"]


def logits_von(ordner: Path, x: np.ndarray, fmt: str) -> np.ndarray:
    if fmt == "onnx":
        import onnxruntime as ort
        s = ort.InferenceSession(str(ordner / "modell.onnx"), providers=["CPUExecutionProvider"])
        return s.run(None, {"x": x.astype(np.float32)})[0]
    import torch
    torch.set_num_threads(4)
    m = torch.jit.load(str(ordner / "modell_torchscript.pt"))
    m.eval()
    with torch.no_grad():
        return m(torch.tensor(x, dtype=torch.float32)).numpy()


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("art", choices=["tabelle", "bilder"])
    ap.add_argument("--ordner", required=True)
    ap.add_argument("--csv")
    ap.add_argument("--npy")
    ap.add_argument("--ausgabe", default="vorhersagen.csv")
    ap.add_argument("--format", default="torchscript", choices=["torchscript", "onnx"])
    a = ap.parse_args()
    ordner = Path(a.ordner)
    if a.art == "tabelle":
        if not a.csv:
            ap.error("--csv fehlt")
        import joblib
        with open(a.csv, newline="", encoding="utf-8") as f:
            zeilen = list(csv.DictReader(f))
        X = np.empty((len(zeilen), len(NUM) + len(KAT)), dtype=object)
        for i, z in enumerate(zeilen):
            for j, c in enumerate(NUM + KAT):
                v = z[c]
                X[i, j] = (np.nan if v == "" else float(v)) if c in NUM else v
        x = joblib.load(ordner / "vorverarbeitung.joblib").transform(X)
    else:
        if not a.npy:
            ap.error("--npy fehlt")
        x = np.load(a.npy)
    lg = logits_von(ordner, x, a.format)
    e = np.exp(lg - lg.max(1, keepdims=True))
    p = e / e.sum(1, keepdims=True)
    klasse = p.argmax(1)
    with open(a.ausgabe, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["zeile", "klasse", "wahrscheinlichkeit"])
        for i, (k, q) in enumerate(zip(klasse, p.max(1))):
            w.writerow([i, int(k), f"{q:.6f}"])
    np.save(str(a.ausgabe) + ".logits.npy", lg)
    return 0


if __name__ == "__main__":
    sys.exit(main())