from typing import Any, Dict, List, Literal, NotRequired, Optional, TypedDict, Union


class ShockEvent(TypedDict):
    type: Literal["shock"]
    t: int
    k: int
    mode: str


class RewireEvent(TypedDict):
    type: Literal["rewire"]
    t: int
    prob: float
    rewired: int


class TopologyActuatorEvent(TypedDict):
    type: Literal["topology_actuator"]
    t: int
    mode: str
    rewired: int
    budget: int
    u: float
    drive: float
    mu_disp_hat: float
    mu_disp_max: float
    low_comm: int
    high_comm: int
    mu_low: float
    mu_high: float
    remove_strategy: str


Event = Union[ShockEvent, RewireEvent, TopologyActuatorEvent]


class TopologyActuatorTimeseriesBase(TypedDict):
    active: bool


class TopologyActuatorTimeseriesActive(TopologyActuatorTimeseriesBase):
    active: Literal[True]
    rewired: int
    budget: int
    drive: float
    mu_disp_hat: float
    mu_disp_max: float
    mode: str


class TopologyActuatorTimeseriesInactive(TopologyActuatorTimeseriesBase):
    active: Literal[False]
    cooldown_remaining: NotRequired[int]


TopologyActuatorTimeseries = Union[TopologyActuatorTimeseriesActive, TopologyActuatorTimeseriesInactive]


class GraphSeries(TypedDict):
    n: int
    m: int


class TruthSeries(TypedDict):
    mu: float
    eta: float
    h: float
    v_mu: float
    v_eta: float
    v: float
    mu_disp: float
    in_e: bool


class ObsSeries(TypedDict):
    cc: float
    cd: float
    dd: float


class DetectorSeries(TypedDict):
    b: float
    risk: float
    b_mu: float
    b_eta: float
    b_v: float
    b_disp: float
    mu_hat: float
    eta_hat: float
    h_hat: float
    mu_disp_hat: float
    v_mu_hat: float
    v_eta_hat: float
    v_hat: float


class ControllerSeries(TypedDict):
    u: float
    u_target: float
    b_eff: float
    risk_eff: float


class DynamicsSeries(TypedDict):
    beta: float
    p_flip: float
    theta: float


class DiagnosticsSeries(TypedDict):
    binding_hat: str
    binding_truth: str
    violated_hat: List[str]
    violated_truth: List[str]
    margins_hat: Dict[str, float]
    margins_truth: Dict[str, float]
    dmu_hat: float
    db_mu: float
    stuck_condition: Optional[bool]
    stuck: Optional[bool]
    stuck_streak: Optional[int]


class MesoSeries(TypedDict):
    k: Optional[int]
    min_total_edges: Optional[float]
    truth_mu_by_comm: Optional[List[float]]


class SeriesRow(TypedDict):
    t: int
    events: List[Event]
    graph: GraphSeries
    topology_actuator: Optional[TopologyActuatorTimeseries]
    truth: TruthSeries
    obs: ObsSeries
    detector: DetectorSeries
    controller: ControllerSeries
    dynamics: DynamicsSeries
    diagnostics: DiagnosticsSeries
    meso: MesoSeries


class PythonProvenance(TypedDict):
    version: str
    platform: str


class Provenance(TypedDict):
    seed: int
    rng_streams: Dict[str, int]
    rng_stream_method: str
    config_path: str
    config_text_sha256: str
    config_canonical_sha256: str
    python: PythonProvenance


class RunOutput(TypedDict):
    provenance: Provenance
    config: Dict[str, Any]
    summary: Dict[str, Any]
    series: List[SeriesRow]