"""The onboard loop and the scenario runner used by every evaluation (dev and held-out)."""
from __future__ import annotations

import random
import zlib

from .sim import FAULTS, Spacecraft

PHYSICAL = {"switch_thruster_branch", "mark_sensor_invalid", "switch_heater_backup", "switch_to_spare_wheel", "switch_attitude_sensor",
            "shed_loads", "safe_mode"}


def scenario(fault, seed, profile="europa_clipper"):
    """a fault class x seed -> injection parameters (onset, zone, wheel, sensor, magnitude), deterministic"""
    # zlib.crc32, not hash(): str hashing is randomised per process, which made scenarios differ between runs
    r = random.Random(zlib.crc32(("%s:%d:%s" % (fault, seed, profile)).encode()))
    zone = {"heater_stuck_on": r.choice(["battery", "tank"]), "heater_stuck_off": r.choice(["tank", "battery"]),
            "sensor_bias": r.choice(["battery", "tank", "avionics"]), "seu_spike": r.choice(["battery", "tank", "avionics"])
            }.get(fault, "battery")
    return {"profile": profile, "fault": fault, "seed": seed, "onset": r.randint(40, 80), "zone": zone, "wheel": r.randint(0, 2),
            "sensor": r.choice("AB"), "magnitude": round(r.random(), 3)}


def run(sc_params, onboard, ticks=320):
    """run one scenario with an onboard decision function; returns the trace the scorer needs"""
    sc = Spacecraft(seed=sc_params["seed"] * 7 + 1, profile=sc_params.get("profile", "europa_clipper"))
    actions, receipts = [], []
    for t in range(ticks):
        if t == sc_params["onset"]:
            sc.inject(sc_params["fault"], zone=sc_params["zone"], wheel=sc_params["wheel"], sensor=sc_params["sensor"],
                      magnitude=sc_params["magnitude"])
        sc.step()
        tm = sc.telemetry()
        r = onboard(tm)
        if r:
            receipts.append(r)
            if r["action"] in PHYSICAL:
                sc.command(r["action"], **r["arg"])
                actions.append((sc.t, r["action"], r["arg"]))
    return {"actions": actions, "receipts": receipts, "damage": sc.damage, "safe": sc.safe, "final_soc": sc.soc}


def score(params, trace):
    from .sim import CORRECT, HARMFUL
    f = params["fault"]
    acts = [a for _, a, _ in trace["actions"]]
    first = trace["actions"][0] if trace["actions"] else None
    correct = CORRECT[f]
    if f in ("nominal", "seu_spike"):
        ok = not acts
    else:
        ok = bool(first) and first[1] in correct and _right_target(params, first)
    harmful = [a for a in acts if a in HARMFUL.get(f, set())] or ([first[1]] if first and not ok and f not in
                                                                  ("nominal", "seu_spike") and first[1] != "safe_mode" else [])
    latency = (first[0] - params["onset"]) if (first and ok) else None
    return {"ok": ok, "harmful": bool(harmful), "first": first[1] if first else None, "latency_min": latency,
            "safe_mode": "safe_mode" in acts, "damage": bool(trace["damage"])}


def _right_target(params, first):
    _, action, arg = first
    if "zone" in arg and params["fault"] in ("heater_stuck_on", "heater_stuck_off", "sensor_bias"):
        if arg["zone"] != params["zone"]:
            return False
    if params["fault"] == "sensor_bias" and arg.get("sensor") not in (None, params["sensor"]):
        return False
    if params["fault"] == "wheel_friction" and arg.get("wheel") not in (None, params["wheel"]):
        return False
    return True


__all__ = ["FAULTS", "scenario", "run", "score"]