"""Autoregressive shower rollout driver.

Steps the two-stage GIANT surrogate forward into a full particle shower: each
primary post-step becomes the next pre-step, secondaries are pushed as new tracks,
and the material/layer_id conditioning at every step comes from a `GeometryOracle`
(the surrogate does not predict them).

Tracks are advanced breadth-first: every sweep steps all currently-active tracks
once (in `batch_size` chunks), so many tracks share each model forward pass. A
track terminates on one of the recorded `termination_reason`s in constants.py.

Energy accounting: on every terminal stop except escape, the track's remaining
energy is deposited locally so the shower conserves energy; escaped energy is
treated as detector leakage and not deposited.
"""

from __future__ import annotations

import numpy as np
import torch

from giant.constants import (
    TERM_ENERGY_CUTOFF,
    TERM_ESCAPED,
    TERM_MAX_STEPS,
    TERM_NATURAL_END,
    TERM_UNKNOWN_PDG,
)
from giant.data.transforms import (
    Normalizer,
    build_cond_features,
    decode_secondaries,
    energy_simplex_decode,
    inv_local_frame_rotation,
    inv_log_transform,
    reconstruct_post_pos,
)
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx

# Record columns produced per step / per terminal marker.
_RECORD_KEYS = [
    "event_id",
    "track_id",
    "parent_id",
    "generation",
    "step_no",
    "pdg",
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "post_x",
    "post_y",
    "post_z",
    "post_E",
    "post_dx",
    "post_dy",
    "post_dz",
    "edep",
    "step_length",
    "material",
    "layer_id",
    "n_sec_pred",
    "termination_reason",
]


def _empty_frontier() -> dict[str, np.ndarray]:
    return {
        "event_id": np.empty(0, dtype=np.int64),
        "track_id": np.empty(0, dtype=np.int64),
        "parent_id": np.empty(0, dtype=np.int64),
        "generation": np.empty(0, dtype=np.int64),
        "step_in_track": np.empty(0, dtype=np.int64),
        "pdg": np.empty(0, dtype=np.int64),
        "pre_pos": np.empty((0, 3), dtype=np.float64),
        "pre_E": np.empty(0, dtype=np.float64),
        "pre_dir": np.empty((0, 3), dtype=np.float64),
    }


def _concat_frontiers(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarray]:
    parts = [p for p in parts if len(p["event_id"]) > 0]
    if not parts:
        return _empty_frontier()
    return {k: np.concatenate([p[k] for p in parts], axis=0) for k in parts[0]}


class _Recorder:
    """Accumulates per-step rows into column lists, materialised at the end."""

    def __init__(self) -> None:
        self._cols: dict[str, list] = {k: [] for k in _RECORD_KEYS}

    def add(self, **cols) -> None:
        n = len(cols["event_id"])
        if n == 0:
            return
        for k in _RECORD_KEYS:
            v = cols[k]
            self._cols[k].append(np.asarray(v).reshape(n))

    def to_dict(self) -> dict[str, np.ndarray]:
        out = {}
        for k, chunks in self._cols.items():
            if chunks:
                out[k] = np.concatenate(chunks, axis=0)
            else:
                out[k] = np.empty(
                    0,
                    dtype=object
                    if k in ("material", "termination_reason")
                    else np.float64,
                )
        return out


def make_seed_frontier(
    event_id: np.ndarray,
    pdg: np.ndarray,
    pre_pos: np.ndarray,
    pre_E: np.ndarray,
    pre_dir: np.ndarray,
) -> tuple[dict[str, np.ndarray], dict[int, int]]:
    """Build the initial frontier from primary entry states.

    Returns (frontier, event_track_count) where the latter tracks the next
    unused track_id per event (each primary gets a fresh id starting from 0).
    """
    event_id = np.asarray(event_id, dtype=np.int64)
    n = len(event_id)
    track_id = np.empty(n, dtype=np.int64)
    counts: dict[int, int] = {}
    for i, ev in enumerate(event_id.tolist()):
        c = counts.get(ev, 0)
        track_id[i] = c
        counts[ev] = c + 1

    dir_ = np.asarray(pre_dir, dtype=np.float64)
    dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None)

    frontier = {
        "event_id": event_id,
        "track_id": track_id,
        "parent_id": np.full(n, -1, dtype=np.int64),
        "generation": np.zeros(n, dtype=np.int64),
        "step_in_track": np.zeros(n, dtype=np.int64),
        "pdg": np.asarray(pdg, dtype=np.int64),
        "pre_pos": np.asarray(pre_pos, dtype=np.float64),
        "pre_E": np.asarray(pre_E, dtype=np.float64),
        "pre_dir": dir_,
    }
    return frontier, counts


def _terminal_rows(tr: dict[str, np.ndarray], sel: np.ndarray, reason: str, edep):
    """Assemble terminal-marker record columns for the selected tracks."""
    pos = tr["pre_pos"][sel]
    dir_ = tr["pre_dir"][sel]
    n = int(sel.sum())
    return dict(
        event_id=tr["event_id"][sel],
        track_id=tr["track_id"][sel],
        parent_id=tr["parent_id"][sel],
        generation=tr["generation"][sel],
        step_no=tr["step_in_track"][sel],
        pdg=tr["pdg"][sel],
        pre_x=pos[:, 0],
        pre_y=pos[:, 1],
        pre_z=pos[:, 2],
        pre_E=tr["pre_E"][sel],
        pre_dx=dir_[:, 0],
        pre_dy=dir_[:, 1],
        pre_dz=dir_[:, 2],
        post_x=pos[:, 0],
        post_y=pos[:, 1],
        post_z=pos[:, 2],
        post_E=np.zeros(n),
        post_dx=dir_[:, 0],
        post_dy=dir_[:, 1],
        post_dz=dir_[:, 2],
        edep=np.asarray(edep, dtype=np.float64).reshape(n),
        step_length=np.zeros(n),
        material=tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
        layer_id=tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
        n_sec_pred=np.zeros(n, dtype=np.int64),
        termination_reason=np.full(n, reason, dtype=object),
    )


@torch.no_grad()
def rollout(
    stage1_model: torch.nn.Module,
    sec_decoder: torch.nn.Module,
    oracle,
    seeds: dict[str, np.ndarray],
    cond_norm: Normalizer,
    tgt_norm: Normalizer,
    pdg_map: dict[int, int],
    mat_map: dict[str, int],
    *,
    energy_cutoff: float,
    max_steps: int,
    steps: int = 10,
    batch_size: int = 4096,
    device: torch.device | None = None,
    max_tracks_per_event: int | None = None,
    escape_threshold: float | None = None,
) -> dict[str, np.ndarray]:
    """Run showers to completion; return a step-record dict (see _RECORD_KEYS)."""
    device = device or torch.device("cpu")
    stage1_model.eval()
    sec_decoder.eval()
    if escape_threshold is not None:
        oracle.escape_threshold = float(escape_threshold)

    pdg_map_inv = {v: k for k, v in pdg_map.items()}
    pdg_emb_weight = stage1_model.pdg_embedding_weight()

    frontier, counts = make_seed_frontier(
        seeds["event_id"],
        seeds["pdg"],
        seeds["pre_pos"],
        seeds["pre_E"],
        seeds["pre_dir"],
    )
    rec = _Recorder()

    while len(frontier["event_id"]) > 0:
        next_parts: list[dict[str, np.ndarray]] = []
        n_total = len(frontier["event_id"])
        for start in range(0, n_total, batch_size):
            chunk = {k: v[start : start + batch_size] for k, v in frontier.items()}
            next_parts.append(
                _step_chunk(
                    chunk,
                    stage1_model,
                    sec_decoder,
                    oracle,
                    cond_norm,
                    tgt_norm,
                    pdg_map,
                    mat_map,
                    pdg_map_inv,
                    pdg_emb_weight,
                    rec,
                    counts,
                    energy_cutoff,
                    max_steps,
                    steps,
                    device,
                    max_tracks_per_event,
                )
            )
        frontier = _concat_frontiers(next_parts)

    return rec.to_dict()


def _step_chunk(
    tr,
    stage1_model,
    sec_decoder,
    oracle,
    cond_norm,
    tgt_norm,
    pdg_map,
    mat_map,
    pdg_map_inv,
    pdg_emb_weight,
    rec,
    counts,
    energy_cutoff,
    max_steps,
    steps,
    device,
    max_tracks_per_event,
) -> dict[str, np.ndarray]:
    """Advance one chunk of tracks by a single step; return the next frontier."""
    n = len(tr["event_id"])

    # --- Geometry lookup + material/layer conditioning ---
    material, layer_id, escaped = oracle.query(tr["pre_pos"])
    tr = dict(tr)
    tr["_material"] = material
    tr["_layer_id"] = layer_id

    known_pdg = np.array([int(p) in pdg_map for p in tr["pdg"]], dtype=bool)

    # --- Pre-step termination gates (in priority order; each track picks one) ---
    stop = np.zeros(n, dtype=bool)
    escaped_sel = escaped & ~stop
    rec.add(
        **_terminal_rows(
            tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))
        )
    )
    stop |= escaped_sel

    unknown_sel = ~known_pdg & ~stop
    rec.add(
        **_terminal_rows(
            tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]
        )
    )
    stop |= unknown_sel

    cutoff_sel = (tr["pre_E"] < energy_cutoff) & ~stop
    rec.add(
        **_terminal_rows(
            tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]
        )
    )
    stop |= cutoff_sel

    maxstep_sel = (tr["step_in_track"] >= max_steps) & ~stop
    rec.add(
        **_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel])
    )
    stop |= maxstep_sel

    active = ~stop
    if not active.any():
        return _empty_frontier()

    tr = {k: v[active] for k, v in tr.items()}
    material = tr["_material"]
    layer_id = tr["_layer_id"]

    # --- Build conditioning and run the two stages ---
    cond_dict = {
        "pre_pos": tr["pre_pos"],
        "pre_E": tr["pre_E"],
        "pre_dir": tr["pre_dir"],
        "layer_id": layer_id,
        "material": material,
        "pdg": tr["pdg"],
    }
    cond_cont, cond_cat = build_cond_features(cond_dict, pdg_map, mat_map, cond_norm)
    cc = torch.from_numpy(cond_cont).float().to(device)
    ck = torch.from_numpy(cond_cat).long().to(device)

    stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps)
    raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())

    step_length = inv_log_transform(raw[:, 0])
    edep, e_sec, post_E, _delta = energy_simplex_decode(raw[:, 1:3], tr["pre_E"])

    post_dir_local = raw[:, 3:6].copy()
    post_dir_local /= np.clip(
        np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None
    )
    post_dir_world = inv_local_frame_rotation(tr["pre_dir"], post_dir_local)

    travel_dir_local = raw[:, 6:9].copy()
    travel_dir_local /= np.clip(
        np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None
    )
    post_pos = reconstruct_post_pos(
        tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local
    )

    n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)

    # --- Secondaries ---
    sec_cont, sec_type_emb, _valid = sample_secondaries(
        sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
    )
    sec_pdg_idx = snap_type_to_pdg_idx(sec_type_emb, pdg_emb_weight)
    sec_E, sec_dir_world, sec_pdg_code, sec_valid = decode_secondaries(
        sec_cont.cpu().numpy(),
        sec_pdg_idx.cpu().numpy(),
        n_sec_np,
        e_sec,
        tr["pre_dir"],
        pdg_map_inv,
    )

    edep = edep.astype(np.float64)
    post_E = post_E.astype(np.float64)

    # --- Spawn secondaries (with per-event track cap) ---
    new_tracks, dropped_edep = _spawn_secondaries(
        tr,
        post_pos,
        sec_valid,
        sec_E,
        sec_dir_world,
        sec_pdg_code,
        counts,
        max_tracks_per_event,
    )
    # Energy bookkeeping so each step conserves exactly (edep + carried + post_E
    # == pre_E): the primary lost `e_sec` to secondaries, but the decoded
    # secondaries only carry `sec_E[valid].sum()`. Deposit the unallocated
    # residual locally, plus the energy of any sub-cap secondaries we dropped.
    sec_E_valid_sum = (sec_E * sec_valid).sum(axis=1)
    residual = np.maximum(e_sec - sec_E_valid_sum, 0.0)
    edep = edep + residual + dropped_edep

    # --- Record the stepped rows; mark natural_end where the primary died ---
    natural = post_E <= 0.0
    reason = np.where(natural, TERM_NATURAL_END, "").astype(object)
    rec.add(
        event_id=tr["event_id"],
        track_id=tr["track_id"],
        parent_id=tr["parent_id"],
        generation=tr["generation"],
        step_no=tr["step_in_track"],
        pdg=tr["pdg"],
        pre_x=tr["pre_pos"][:, 0],
        pre_y=tr["pre_pos"][:, 1],
        pre_z=tr["pre_pos"][:, 2],
        pre_E=tr["pre_E"],
        pre_dx=tr["pre_dir"][:, 0],
        pre_dy=tr["pre_dir"][:, 1],
        pre_dz=tr["pre_dir"][:, 2],
        post_x=post_pos[:, 0],
        post_y=post_pos[:, 1],
        post_z=post_pos[:, 2],
        post_E=post_E,
        post_dx=post_dir_world[:, 0],
        post_dy=post_dir_world[:, 1],
        post_dz=post_dir_world[:, 2],
        edep=edep,
        step_length=step_length,
        material=material,
        layer_id=layer_id,
        n_sec_pred=n_sec_np,
        termination_reason=reason,
    )

    # --- Continue surviving primaries ---
    cont = ~natural
    cont_frontier = {
        "event_id": tr["event_id"][cont],
        "track_id": tr["track_id"][cont],
        "parent_id": tr["parent_id"][cont],
        "generation": tr["generation"][cont],
        "step_in_track": tr["step_in_track"][cont] + 1,
        "pdg": tr["pdg"][cont],
        "pre_pos": post_pos[cont],
        "pre_E": post_E[cont],
        "pre_dir": post_dir_world[cont],
    }
    return _concat_frontiers([cont_frontier, new_tracks])


def _spawn_secondaries(
    tr,
    post_pos,
    sec_valid,
    sec_E,
    sec_dir_world,
    sec_pdg_code,
    counts,
    max_tracks_per_event,
) -> tuple[dict[str, np.ndarray], np.ndarray]:
    """Turn valid secondaries into new tracks; return (frontier, per-parent dropped edep).

    Secondaries are born at their parent's post_pos. When `max_tracks_per_event`
    is set and an event is at its cap, further secondaries are not spawned; their
    energy is returned as `dropped_edep` (indexed by parent row) so it is
    deposited into the parent step instead of vanishing.
    """
    B = len(tr["event_id"])
    dropped_edep = np.zeros(B, dtype=np.float64)

    pr, sl = np.nonzero(sec_valid)  # parent-row idx, slot idx
    if len(pr) == 0:
        return _empty_frontier(), dropped_edep

    # Assign a fresh per-event track_id to each candidate in stable parent order,
    # applying the per-event cap. The candidate count per chunk is small
    # (<= batch_size * K_MAX), so a plain loop is clear and fast enough.
    order = np.lexsort((sl, pr))  # group by parent row, slot ascending
    kept_pr, kept_sl, kept_tid = [], [], []
    for j in order:
        ev = int(tr["event_id"][pr[j]])
        cur = counts.get(ev, 0)
        if max_tracks_per_event is not None and cur >= max_tracks_per_event:
            dropped_edep[pr[j]] += float(sec_E[pr[j], sl[j]])
            continue
        kept_pr.append(pr[j])
        kept_sl.append(sl[j])
        kept_tid.append(cur)
        counts[ev] = cur + 1

    if not kept_pr:
        return _empty_frontier(), dropped_edep

    pr_k = np.array(kept_pr, dtype=np.int64)
    sl_k = np.array(kept_sl, dtype=np.int64)
    tid_k = np.array(kept_tid, dtype=np.int64)

    frontier = {
        "event_id": tr["event_id"][pr_k],
        "track_id": tid_k,
        "parent_id": tr["track_id"][pr_k],
        "generation": tr["generation"][pr_k] + 1,
        "step_in_track": np.zeros(len(pr_k), dtype=np.int64),
        "pdg": sec_pdg_code[pr_k, sl_k].astype(np.int64),
        "pre_pos": post_pos[pr_k],
        "pre_E": sec_E[pr_k, sl_k].astype(np.float64),
        "pre_dir": sec_dir_world[pr_k, sl_k].astype(np.float64),
    }
    return frontier, dropped_edep
