import json, math, random
from dataclasses import dataclass
from pathlib import Path
import numpy as np

@dataclass(frozen=True)
class Transition:
    name: str
    inp: int
    out: int

class PetriNet:
    def __init__(self, n_places, transitions, initial, admissible):
        self.n_places = n_places
        self.transitions = transitions
        self.initial = tuple(initial)
        self.admissible = {tuple(m) for m in admissible}

    def enabled(self, t, m):
        return m[t.inp] >= 1

    def fire(self, t, m):
        x = list(m); x[t.inp] -= 1; x[t.out] += 1
        return tuple(x)

    def allowed(self, m):
        return [t for t in self.transitions
                if self.enabled(t, m) and self.fire(t, m) in self.admissible]


def make_rover():
    # One-hot marking: position x battery b, plus one unsafe sink.
    states = [(x, b) for x in range(5) for b in range(3)]
    place = {s:i for i, s in enumerate(states)}
    sink = len(states)
    actions = ["left", "right", "recharge", "wait"]
    transitions = []
    for s in states:
        x, b = s
        for a in actions:
            valid = True
            if a == "left":
                valid = x > 0 and b > 0
                ns = (x-1, b-1) if valid else None
            elif a == "right":
                valid = x < 4 and b > 0
                ns = (x+1, b-1) if valid else None
            elif a == "recharge":
                valid = x == 0 and b < 2
                ns = (x, 2) if valid else None
            else:
                ns = s
            transitions.append(Transition(f"{s}:{a}", place[s], place[ns] if valid else sink))
    init = tuple(1 if i == place[(0, 2)] else 0 for i in range(sink+1))
    # M_obs is the exact reachable set under safe transitions, not merely all bit patterns.
    seen = {init}; frontier = [init]
    temp = PetriNet(sink+1, transitions, init, seen)
    while frontier:
        m = frontier.pop()
        for t in temp.transitions:
            if temp.enabled(t, m):
                q = temp.fire(t, m)
                if q[sink] == 0 and q not in seen:
                    seen.add(q); frontier.append(q)
    return PetriNet(sink+1, transitions, init, seen), states, actions, place, sink


def argmax_action(logits):
    return int(np.argmax(logits))


def run_policy(net, states, actions, place, sink, shield, rng, episodes=10000):
    # Fixed random linear neural policy: same scores are used by both controllers.
    d = len(states)
    W = rng.normal(0, 1, size=(len(actions), 2))
    violations = rejects = deadlocks = steps = successes = 0
    for _ in range(episodes):
        m = net.initial
        for _step in range(40):
            idx = int(np.argmax(np.asarray(m[:-1])))
            x, b = states[idx]
            logits = W @ np.array([x / 4.0, b / 2.0])
            candidates = [t for t in net.transitions if t.inp == idx]
            # actions are ordered consistently by construction.
            if shield:
                allowed = [t for t in candidates if net.fire(t, m) in net.admissible]
                if not allowed:
                    deadlocks += 1; break
                best = int(np.argmax([logits[actions.index(t.name.split(':')[1])] for t in allowed]))
                t = allowed[best]
                if int(np.argmax(logits)) != actions.index(t.name.split(':')[1]): rejects += 1
            else:
                ai = argmax_action(logits); t = candidates[ai]
                if net.fire(t, m)[sink] == 1:
                    violations += 1; break
            m = net.fire(t, m); steps += 1
            # Reaching x=4 is a simple task success signal.
            if states[int(np.argmax(np.asarray(m[:-1])))][0] == 4:
                successes += 1; break
    return {"violations": violations, "rejections": rejects, "deadlocks": deadlocks,
            "steps": steps, "successes": successes, "episodes": episodes}


def rejection_sweep():
    # Prediction: with equal logits and uniform sampling, rejection probability = 1-k/n.
    rows = []
    rng = np.random.default_rng(123)
    for n in [2, 4, 8, 16]:
        for k in range(1, n):
            trials = 20000
            chosen = rng.integers(0, n, size=trials)
            observed = float(np.mean(chosen >= k))
            predicted = 1.0 - k / n
            rows.append({"n": n, "allowed": k, "predicted": predicted, "observed": observed,
                         "abs_error": abs(observed-predicted)})
    return rows


def main():
    random.seed(7); np.random.seed(7)
    net, states, actions, place, sink = make_rover()
    # Core math checks: every accepted successor is in M_obs; BFS is finite and nonempty.
    closure_bad = sum(net.fire(t, m) not in net.admissible
                      for m in net.admissible for t in net.allowed(m))
    # Compare exact same policy logits with and without the transition mask.
    base = run_policy(net, states, actions, place, sink, False, np.random.default_rng(99))
    shield = run_policy(net, states, actions, place, sink, True, np.random.default_rng(99))
    sweep = rejection_sweep()
    max_err = max(r["abs_error"] for r in sweep)
    result = {
      "petri_net": {"places": net.n_places, "transitions": len(net.transitions),
                    "admissible_markings": len(net.admissible), "closure_bad": closure_bad},
      "predictions": {
        "inductive_closure": {"predicted": "0 unsafe accepted successors", "observed": closure_bad},
        "uniform_rejection": {"predicted": "P(reject)=1-k/n", "max_abs_error": max_err,
                               "rows": sweep},
        "shield_violations": {"predicted": 0, "observed": shield["violations"]}
      },
      "controller_comparison": {"baseline_unshielded": base, "idea_shielded": shield}
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()