import numpy as np

_EPS = 1e-8

# Floor added to each energy fraction before taking log-ratios so the simplex
# coordinates stay finite on the heavily populated boundary: e_sec is 0 in ~half
# of all steps and post_E is 0 at every track end. Kept tiny (1e-5 of pre_E) so
# the conservation it slightly softens is physically negligible (~0.001%).
_SIMPLEX_FLOOR = 1e-5


def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
    return np.log(np.asarray(x, dtype=np.float32) + eps)


def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray:
    return np.exp(np.asarray(y, dtype=np.float32)) - eps


def energy_simplex_encode(
    edep: np.ndarray,
    e_sec: np.ndarray,
    post_E: np.ndarray,
    pre_E: np.ndarray,
) -> np.ndarray:
    """Encode (edep, e_sec, post_E) as 2 additive-log-ratio (ALR) coordinates.

    The three energies are first expressed as fractions of pre_E that sum to 1:
    post_E is kept exact (so the particle's retained energy — the next step's
    input during a rollout — is preserved), and the energy actually lost,
    `delta_e = pre_E - post_E`, is split between local deposit and secondaries
    in the recorded edep:e_sec ratio (re-attributing any sub-threshold / rest-mass
    leakage proportionally so the three fractions sum to exactly 1). A small floor
    (`_SIMPLEX_FLOOR`) keeps the log-ratios finite where a fraction is 0.

    Returns an (N, 2) array of ALR coordinates referenced to the post fraction;
    `energy_simplex_decode` is its inverse (up to the floor softening).
    """
    pre_E = np.maximum(np.asarray(pre_E, dtype=np.float32), _EPS)
    post_E = np.clip(np.asarray(post_E, dtype=np.float32), 0.0, pre_E)
    delta_e = pre_E - post_E
    lost = np.asarray(edep, dtype=np.float32) + np.asarray(e_sec, dtype=np.float32)
    has_loss = lost > _EPS
    scale = np.where(has_loss, delta_e / np.maximum(lost, _EPS), 0.0)
    # Where nothing was recorded as deposited/secondary but energy was lost,
    # attribute all of delta_e to local deposit.
    edep_r = np.where(has_loss, edep * scale, delta_e)
    e_sec_r = np.where(has_loss, e_sec * scale, 0.0)

    f = np.stack([edep_r, e_sec_r, post_E], axis=-1) / pre_E[:, None]  # (N,3), sums≈1
    f = (f + _SIMPLEX_FLOOR) / (1.0 + 3.0 * _SIMPLEX_FLOOR)
    log_f = np.log(f)
    z = log_f[:, :2] - log_f[:, 2:3]  # ALR vs post reference
    return z.astype(np.float32)


def energy_simplex_decode(
    z: np.ndarray, pre_E: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Inverse of `energy_simplex_encode`: ALR coords + pre_E → physical energies.

    A softmax over `[z_edep, z_sec, 0]` recovers the three simplex fractions, so
    `edep + e_sec + post_E == pre_E` holds by construction (the energy-conservation
    inductive bias). Returns `(edep, e_sec, post_E, delta_e)` in physical units.
    """
    z = np.asarray(z, dtype=np.float32)
    pre_E = np.asarray(pre_E, dtype=np.float32)
    logits = np.concatenate([z, np.zeros((len(z), 1), dtype=np.float32)], axis=1)
    logits = logits - logits.max(axis=1, keepdims=True)
    f = np.exp(logits)
    f /= f.sum(axis=1, keepdims=True)  # sums to 1 exactly → exact conservation
    E = f * pre_E[:, None]
    edep, e_sec, post_E = E[:, 0], E[:, 1], E[:, 2]
    delta_e = pre_E - post_E
    return (
        edep.astype(np.float32),
        e_sec.astype(np.float32),
        post_E.astype(np.float32),
        delta_e.astype(np.float32),
    )


def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
    """Unit rotation axis `pre_dir × ẑ`, closed-form since ẑ = [0, 0, 1] is constant.

    `cross(a, [0,0,1]) = [a_y, -a_x, 0]` — substituting the constant operand
    avoids a generic `np.cross` call (shape/broadcast handling for an
    arbitrary second operand) on every row; profiling on a 114M-row file
    showed `np.cross` as the single hottest call inside this rotation.
    """
    axis = np.stack(
        [pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
    )
    axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)  # (N,1)
    # Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
    # are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
    safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
    return np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)


def _cross_with_z_axis(axis: np.ndarray, v: np.ndarray) -> np.ndarray:
    """`axis × v`, closed-form since `axis` from `_rodrigues_axis` always has z = 0.

    `cross([ax,ay,0], [bx,by,bz]) = [ay*bz, -ax*bz, ax*by - ay*bx]`.
    """
    ax, ay = axis[:, 0:1], axis[:, 1:2]
    bx, by, bz = v[:, 0:1], v[:, 1:2], v[:, 2:3]
    return np.concatenate([ay * bz, -ax * bz, ax * by - ay * bx], axis=1)


def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarray:
    """Rotate post_dir into the local frame where pre_dir maps to ẑ (Rodrigues).

    Preserves the angle between pre_dir and post_dir; the result has post_dir
    expressed relative to a coordinate system in which the incoming particle
    travels along +z.
    """
    cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0)  # (N,1); dot with ẑ = z-component
    sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2))  # (N,1)

    axis = _rodrigues_axis(pre_dir)  # (N,3); zero-z, zero-norm when pre_dir ∥ ẑ
    kxv = _cross_with_z_axis(axis, post_dir)  # (N,3)
    kdv = (axis * post_dir).sum(axis=1, keepdims=True)  # (N,1)

    return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
        np.float32
    )


class Normalizer:
    def __init__(self) -> None:
        self.mean: np.ndarray | None = None
        self.std: np.ndarray | None = None

    def fit(self, X: np.ndarray) -> "Normalizer":
        self.mean = X.mean(axis=0).astype(np.float32)
        self.std = X.std(axis=0).astype(np.float32)
        self.std = np.where(self.std < _EPS, 1.0, self.std).astype(np.float32)
        return self

    def transform(self, X: np.ndarray) -> np.ndarray:
        return ((X - self.mean) / self.std).astype(np.float32)

    def inverse_transform(self, X: np.ndarray) -> np.ndarray:
        return (X * self.std + self.mean).astype(np.float32)

    def to_dict(self) -> dict:
        assert self.mean is not None and self.std is not None, "Normalizer not fitted"
        return {"mean": self.mean.tolist(), "std": self.std.tolist()}

    @classmethod
    def from_dict(cls, d: dict) -> "Normalizer":
        obj = cls()
        obj.mean = np.array(d["mean"], dtype=np.float32)
        obj.std = np.array(d["std"], dtype=np.float32)
        return obj


class _WelfordAccumulator:
    """Streaming mean/variance (Welford's online algorithm, batch update).

    Use to fit a Normalizer over data that doesn't fit in memory:
        acc = _WelfordAccumulator(n_features)
        for chunk in data:
            acc.update(chunk)
        normalizer = acc.to_normalizer()
    """

    def __init__(self, n_features: int) -> None:
        self.n = 0
        self._mean = np.zeros(n_features, dtype=np.float64)
        self._M2 = np.zeros(n_features, dtype=np.float64)

    def update(self, X: np.ndarray) -> None:
        X = np.asarray(X, dtype=np.float64)
        B = X.shape[0]
        new_n = self.n + B
        delta = X - self._mean
        self._mean += delta.sum(0) / new_n
        delta2 = X - self._mean
        self._M2 += (delta * delta2).sum(0)
        self.n = new_n

    def to_normalizer(self) -> "Normalizer":
        norm = Normalizer()
        norm.mean = self._mean.astype(np.float32)
        std = np.sqrt(self._M2 / max(self.n, 1)).astype(np.float32)
        norm.std = np.where(std < _EPS, 1.0, std).astype(np.float32)
        return norm


def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
    """World-frame unit vector pointing from pre_pos to post_pos.

    Kept independent of `step_length`: that scalar already encodes the
    magnitude of this displacement, so this function only ever returns
    direction (norm-guarded the same way as `local_frame_rotation`'s axis).
    """
    disp = post_pos - pre_pos
    norm = np.linalg.norm(disp, axis=1, keepdims=True)
    safe_norm = np.where(norm < 1e-7, 1.0, norm)
    return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(
        np.float32
    )


def reconstruct_post_pos(
    pre_pos: np.ndarray,
    pre_dir: np.ndarray,
    step_length: np.ndarray,
    travel_dir_local: np.ndarray,
) -> np.ndarray:
    """Inverse of the travel_direction/local_frame_rotation encoding.

    Single source of truth for combining the magnitude (`step_length`) and
    direction (`travel_dir_local`) back into a world-frame post_pos, so
    `step_length` and post_pos stay consistent by construction.
    """
    travel_dir_world = inv_local_frame_rotation(pre_dir, travel_dir_local)
    return (pre_pos + step_length.reshape(-1, 1) * travel_dir_world).astype(np.float32)


def inv_local_frame_rotation(
    pre_dir: np.ndarray, post_dir_local: np.ndarray
) -> np.ndarray:
    """Inverse of local_frame_rotation: rotate from local frame back to world frame.

    Applies R^T (same axis, negative angle) to post_dir_local.
    """
    cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0)  # dot with ẑ = z-component
    sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2))

    axis = _rodrigues_axis(pre_dir)
    kxv = _cross_with_z_axis(axis, post_dir_local)
    kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)

    # Negative angle: sin_t → -sin_t
    return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
        np.float32
    )


def build_cond_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
    mat_map: dict[str, int],
    cond_normalizer: "Normalizer | None" = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Build conditioning arrays only — no target, no post-step variables."""
    cond_cont = np.column_stack(
        [
            data["pre_pos"],
            log_transform(data["pre_E"]),
            data["pre_dir"],
            data["layer_id"].astype(np.float32),
            data["n_sec"].astype(np.float32),
            log_transform(data["e_sec"]),
        ]
    ).astype(np.float32)

    pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
    mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
    cond_cat = np.column_stack([pdg_idx, mat_idx])

    if cond_normalizer is not None:
        cond_cont = cond_normalizer.transform(cond_cont)

    return cond_cont, cond_cat


def build_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
    mat_map: dict[str, int],
    cond_normalizer: Normalizer | None = None,
    target_normalizer: Normalizer | None = None,
    fit: bool = False,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, Normalizer | None, Normalizer | None]:
    """Assemble (cond_cont, cond_cat, target) arrays ready for StepsDataset.

    When fit=True, new Normalizers are fitted on the supplied arrays.
    """
    post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
    travel_dir_local = local_frame_rotation(
        data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
    )

    energy_z = energy_simplex_encode(
        data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
    )  # (N, 2): ALR coords of the deposit/secondary/post energy simplex

    target = np.column_stack(
        [
            log_transform(data["step_length"]),
            energy_z,
            post_dir_local,
            travel_dir_local,
        ]
    ).astype(np.float32)  # (N, 9)

    cond_cont = np.column_stack(
        [
            data["pre_pos"],
            log_transform(data["pre_E"]),
            data["pre_dir"],
            data["layer_id"].astype(np.float32),
            data["n_sec"].astype(np.float32),
            log_transform(data["e_sec"]),
        ]
    ).astype(np.float32)  # (N, COND_DIM)

    pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
    mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
    cond_cat = np.column_stack([pdg_idx, mat_idx])  # (N, 2)

    if fit:
        cond_normalizer = Normalizer().fit(cond_cont)
        target_normalizer = Normalizer().fit(target)

    if cond_normalizer is not None:
        cond_cont = cond_normalizer.transform(cond_cont)
    if target_normalizer is not None:
        target = target_normalizer.transform(target)

    return cond_cont, cond_cat, target, cond_normalizer, target_normalizer
