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


@dataclass
class EventCell:
    theta: float = 1.0
    reset: float = 0.0
    leak: float = 1.0

    def run(self, x, events):
        """Process (normalized_time, signed_mass) events in stable time order."""
        u = self.leak * float(x)
        spikes = []
        # Python's stable sort preserves input order for tied timestamps.
        for tau, delta, *rest in sorted(events, key=lambda e: e[0]):
            u += float(delta)
            fired = int(u >= self.theta)
            spikes.append(fired)
            if fired:
                u = self.reset
        return int(any(spikes)), u, spikes

    def aggregate(self, x, events):
        """Baseline: one signed current and one threshold/reset per macro-step."""
        u = self.leak * float(x) + sum(float(e[1]) for e in events)
        fired = int(u >= self.theta)
        if fired:
            u = self.reset
        return fired, u


def order_condition(x, a, b, theta=1.0):
    return (x + a - b < theta) and (x + a >= theta)


def analytic_fraction(a, b, theta=1.0, x_low=0.0, x_high=1.0):
    # x must lie in [theta-a, theta-a+b), intersected with the sampling interval.
    lo = max(x_low, theta - a)
    hi = min(x_high, theta - a + b)
    return max(0.0, hi - lo) / (x_high - x_low)


def toy_verification(rng):
    theta = 1.0
    n = 500_000
    x = rng.uniform(0.0, theta, n)

    # Prediction 1: unsafe/order-sensitive fraction is the interval length.
    a = 0.65
    b_grid = np.array([0.05, 0.15, 0.30, 0.50, 0.80])
    p_obs, p_pred = [], []
    for b in b_grid:
        p_obs.append(np.mean((x + a - b < theta) & (x + a >= theta)))
        p_pred.append(analytic_fraction(a, b, theta))

    # Prediction 2: for fixed b, the fraction grows linearly with a until
    # clipping at the available x interval.
    b = 0.30
    a_grid = np.array([0.10, 0.30, 0.50, 0.70, 0.90, 1.10])
    pa_obs, pa_pred = [], []
    for aa in a_grid:
        pa_obs.append(np.mean((x + aa - b < theta) & (x + aa >= theta)))
        pa_pred.append(analytic_fraction(aa, b, theta))

    # Prediction 3: conditioned on the strict margin interval, random E/I
    # ordering gives a 1/2 expected firing discrepancy.
    m = 0.10
    x0 = theta - a + b / 2.0  # strictly inside condition for all listed b>=.15
    b0 = 0.50
    # Explicitly use many random orderings of two events.
    trials = 200_000
    fires = []
    cell = EventCell(theta=theta)
    for _ in range(trials):
        if rng.random() < 0.5:
            ev = [(0.1, a), (0.9, -b0)]
        else:
            ev = [(0.1, -b0), (0.9, a)]
        fires.append(cell.run(x0, ev)[0])
    random_order_firing = float(np.mean(fires))
    # E-first is one, I-first is zero; baseline is always zero here.
    discrepancy_from_half = abs(random_order_firing - 0.5)

    return {
        "prediction_1_interval_fraction": {
            "parameter": "b", "a": a,
            "rows": [{"b": float(bb), "predicted": float(pp), "observed": float(po),
                      "abs_error": float(abs(pp-po))}
                     for bb, pp, po in zip(b_grid, p_pred, p_obs)]
        },
        "prediction_2_interval_scaling": {
            "parameter": "a", "b": b,
            "rows": [{"a": float(aa), "predicted": float(pp), "observed": float(po),
                      "abs_error": float(abs(pp-po))}
                     for aa, pp, po in zip(a_grid, pa_pred, pa_obs)]
        },
        "prediction_3_random_order_half": {
            "predicted": 0.5, "observed": random_order_firing,
            "abs_error": discrepancy_from_half, "x": x0, "a": a, "b": b0,
            "margin": min(theta-(x0+a-b0), x0+a-theta)
        }
    }


def mini_experiment(rng):
    # Each example has one excitatory and one inhibitory arrival. The reference
    # is causal processing; the baseline collapses both into one signed pulse.
    cell = EventCell(theta=1.0, reset=0.0)
    n = 100_000
    x = rng.uniform(0.0, 1.0, n)
    a = rng.uniform(0.05, 0.95, n)
    b = rng.uniform(0.05, 0.95, n)
    # Generate both causal orderings, equally often.
    order = rng.integers(0, 2, n)
    micro = np.empty(n, dtype=np.int8)
    base = np.empty(n, dtype=np.int8)
    unsafe = np.zeros(n, dtype=bool)
    margins = np.full(n, np.nan)
    for i in range(n):
        if order[i] == 0:  # E first
            events = [(0.1, a[i]), (0.9, -b[i])]
        else:
            events = [(0.1, -b[i]), (0.9, a[i])]
        micro[i] = cell.run(x[i], events)[0]
        base[i] = cell.aggregate(x[i], events)[0]
        unsafe[i] = order_condition(x[i], a[i], b[i])
        if unsafe[i]:
            margins[i] = min(1.0-(x[i]+a[i]-b[i]), x[i]+a[i]-1.0)

    # In this setup aggregate is exactly the net-current model; compare it to
    # the causal reference and expose the predicted 1/2 split on unsafe cases.
    disagreement = float(np.mean(micro != base))
    unsafe_rate = float(np.mean(unsafe))
    unsafe_disagreement = float(np.mean((micro != base)[unsafe])) if unsafe.any() else 0.0
    micro_rate = float(np.mean(micro))
    base_rate = float(np.mean(base))
    return {
        "n": n,
        "unsafe_fraction": unsafe_rate,
        "mean_margin_on_unsafe": float(np.nanmean(margins)),
        "baseline_spike_rate": base_rate,
        "micro_event_spike_rate": micro_rate,
        "overall_disagreement": disagreement,
        "disagreement_conditioned_on_unsafe": unsafe_disagreement,
        "predicted_conditioned_disagreement": 0.5,
        "order_balance": float(np.mean(order == 0))
    }


def main():
    rng = np.random.default_rng(2411)
    toy = toy_verification(rng)
    mini = mini_experiment(rng)
    result = {"seed": 2411, "toy_verification": toy, "mini_experiment": mini}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()