"""
Multivariate Hawkes process calibrator.

Models self-exciting and cross-exciting point processes for
limit order book events: trades, cancellations, spread widenings.

A type-j event at time t_k excites the intensity of type-i events:
  λ_i(t) = μ_i + Σ_j α_ij * Σ_{t_k < t} exp(-β * (t - t_k))

Key applications:
  - Queue depletion probability (trade → more trades)
  - Cancel cascade detection (cancel → more cancels)
  - Spread widening prediction (trade → spread widening)
  - Toxicity anticipation (flow → adverse selection)

Calibration: maximum likelihood via gradient descent on synthetic or
real event data. Braning ratio Σ_j α_ij / β < 1 for stationarity.
"""

from __future__ import annotations

import numpy as np


class HawkesCalibrator:
    """Multivariate Hawkes process with MLE calibration."""

    def __init__(self, n_dimensions: int = 3, beta: float = 2.0):
        self._n_dim = n_dimensions
        self._mu = np.ones(n_dimensions) * 0.5     # baseline intensity
        self._alpha = np.eye(n_dimensions) * 0.1    # excitation matrix
        self._beta = beta                            # decay rate

    # ── Properties ───────────────────────────────────────────

    @property
    def n_dim(self) -> int:
        return self._n_dim

    @property
    def mu(self) -> np.ndarray:
        return self._mu.copy()

    @mu.setter
    def mu(self, value: np.ndarray):
        self._mu = np.asarray(value, dtype=float)

    @property
    def alpha(self) -> np.ndarray:
        return self._alpha.copy()

    @alpha.setter
    def alpha(self, value: np.ndarray):
        self._alpha = np.asarray(value, dtype=float)

    @property
    def beta(self) -> float:
        return self._beta

    @beta.setter
    def beta(self, value: float):
        self._beta = float(value)

    # ── Calibration ──────────────────────────────────────────

    def calibrate(
        self,
        events: list[tuple[int, float]],
        max_time: float,
        learning_rate: float = 0.01,
        iterations: int = 200,
        regularization: float = 0.001,
    ):
        """Calibrate parameters via stochastic gradient descent on log-likelihood.

        Args:
            events: list of (type, timestamp) tuples, must be sorted by time
            max_time: total observation window
            learning_rate: SGD step size
            iterations: number of gradient steps
            regularization: L2 penalty on alpha and mu
        """
        events_by_type = [[] for _ in range(self._n_dim)]
        for ev_type, ev_time in events:
            if 0 <= ev_type < self._n_dim:
                events_by_type[ev_type].append(ev_time)

        for _ in range(iterations):
            grad_mu, grad_alpha = _compute_gradient(
                events_by_type, self._mu, self._alpha, self._beta, max_time
            )
            # Gradient ascent with L2 regularization
            self._mu += learning_rate * (grad_mu - regularization * self._mu)
            self._alpha += learning_rate * (grad_alpha - regularization * self._alpha)
            # Project to valid range
            self._mu = np.maximum(self._mu, 0.001)
            self._alpha = np.maximum(self._alpha, 0.0)
            # Enforce stationarity: branching ratio < 0.99
            row_sums = self._alpha.sum(axis=1)
            for i in range(self._n_dim):
                if row_sums[i] > self._beta * 0.99:
                    self._alpha[i] *= (self._beta * 0.99) / row_sums[i]

    # ── Intensity ────────────────────────────────────────────

    def intensity(
        self,
        dim: int,
        event_history: list[tuple[int, float]],
        current_time: float,
    ) -> float:
        """Compute intensity λ_i(t) given event history."""
        full = hawkes_intensity(
            self._mu, self._alpha, self._beta,
            event_history, current_time, dim=int(dim),
        )
        return float(full)

    # ── Metrics ──────────────────────────────────────────────

    def branching_ratio(self) -> float:
        """Average branching ratio = max_i(Σ_j α_ij / β). Must be < 1."""
        ratios = self._alpha.sum(axis=1) / self._beta
        return float(np.max(ratios))

    def forecast_activity(
        self,
        events: list[tuple[int, float]],
        current_time: float,
        horizon: float = 1.0,
        n_samples: int = 100,
    ) -> np.ndarray:
        """Forecast expected event count per type in the next horizon.

        Uses the branching structure: E[N_i] = μ_i * horizon + Σ_j α_ij/β * current_excitation.
        """
        expected = np.zeros(self._n_dim)
        contribution = np.zeros(self._n_dim)

        # Baseline contribution
        expected += self._mu * horizon

        # Excitation from past events
        for ev_type, ev_time in events:
            if ev_time >= current_time:
                continue
            decay = np.exp(-self._beta * (current_time - ev_time))
            remaining = (1.0 - np.exp(-self._beta * horizon)) / self._beta
            for i in range(self._n_dim):
                contribution[i] += self._alpha[i, ev_type] * decay * remaining

        result = expected + contribution
        return np.maximum(result, 0.0)


# ── Pure functions ──────────────────────────────────────────

def hawkes_intensity(
    mu: np.ndarray,
    alpha: np.ndarray,
    beta: float,
    event_history: list[tuple[int, float]],
    current_time: float,
    dim: int | None = None,
) -> np.ndarray:
    """Vectorized Hawkes intensity.

    λ_i(t) = μ_i + Σ_j α_ij * Σ_{t_k < t} exp(-β(t - t_k))

    If dim is specified, returns scalar for that dimension.
    """
    n_dim = len(mu)
    intensity = mu.copy().astype(float)

    for ev_type, ev_time in event_history:
        if ev_time >= current_time or ev_type >= n_dim:
            continue
        decay = np.exp(-beta * (current_time - ev_time))
        intensity += alpha[:, ev_type] * decay

    if dim is not None:
        return intensity[dim]
    return intensity


def hawkes_log_likelihood(
    events: list[tuple[int, float]],
    mu: np.ndarray,
    alpha: np.ndarray,
    beta: float,
    max_time: float,
) -> float:
    """Compute log-likelihood of observing these events under given parameters."""
    n_dim = len(mu)
    events_by_type = [[] for _ in range(n_dim)]
    for ev_type, ev_time in events:
        if 0 <= ev_type < n_dim:
            events_by_type[ev_type].append(ev_time)

    # Term 1: sum over events log(λ_i(t_k))
    event_ll = 0.0
    for ev_type, ev_time in events:
        lam = mu[ev_type]
        for past_type, past_time in events:
            if past_time >= ev_time:
                break
            lam += alpha[ev_type, past_type] * np.exp(-beta * (ev_time - past_time))
        if lam > 0:
            event_ll += np.log(lam)

    # Term 2: -∫ λ(t) dt (compensator)
    integral = max_time * mu.sum()

    for i in range(n_dim):
        for j in range(n_dim):
            for t_j in events_by_type[j]:
                integral += alpha[i, j] * (1.0 - np.exp(-beta * (max_time - t_j))) / beta

    return event_ll - integral


def generate_hawkes_events(
    mu: np.ndarray,
    alpha: np.ndarray,
    beta: float,
    max_time: float,
    seed: int | None = None,
) -> list[tuple[int, float]]:
    """Generate synthetic events from a multivariate Hawkes process (Ogata thinning).

    Returns list of (type, timestamp) sorted by time.
    """
    rng = np.random.RandomState(seed) if seed is not None else np.random
    n_dim = len(mu)

    # Upper bound for total intensity
    lambda_bar = mu.sum() * 1.5  # conservative upper bound

    events: list[tuple[int, float]] = []
    t = 0.0

    while t < max_time:
        # Generate candidate via Poisson with rate lambda_bar (thinning)
        t += rng.exponential(1.0 / lambda_bar) if lambda_bar > 0 else max_time

        if t >= max_time:
            break

        # Accept with probability λ(t) / lambda_bar
        current_intensity = hawkes_intensity(mu, alpha, beta, events, t)
        total_intensity = current_intensity.sum()

        if total_intensity / lambda_bar > rng.uniform(0, 1):
            # Accept: determine event type proportional to intensity
            probs = current_intensity / total_intensity
            ev_type = rng.choice(n_dim, p=probs)
            events.append((int(ev_type), t))

    return events


# ── Gradient computation ───────────────────────────────────

def _compute_gradient(
    events_by_type: list[list[float]],
    mu: np.ndarray,
    alpha: np.ndarray,
    beta: float,
    max_time: float,
) -> tuple[np.ndarray, np.ndarray]:
    """Compute gradient of log-likelihood wrt mu and alpha."""
    n_dim = len(mu)
    grad_mu = np.zeros(n_dim)
    grad_alpha = np.zeros((n_dim, n_dim))

    # Build flat event list for gradient computation
    all_events = []
    for i, times in enumerate(events_by_type):
        for t in times:
            all_events.append((i, t))
    all_events.sort(key=lambda x: x[1])

    # Gradient of log-likelihood
    for ev_type, ev_time in all_events:
        lam = mu[ev_type]
        past_contributions = {}
        for pt, ptime in all_events:
            if ptime >= ev_time:
                break
            contrib = alpha[ev_type, pt] * np.exp(-beta * (ev_time - ptime))
            lam += contrib
            past_contributions[pt] = contrib

        if lam > 0:
            grad_mu[ev_type] += 1.0 / lam

    # Gradient of alpha: Σ_{t_k} α_{ij} * exp(-β (t - t_k)) contribution
    for i in range(n_dim):
        for j in range(n_dim):
            for t_j in events_by_type[j]:
                decay_integral = (1.0 - np.exp(-beta * (max_time - t_j))) / beta
                grad_alpha[i, j] -= decay_integral

    # Add per-event alpha gradient
    for ev_type, ev_time in all_events:
        lam = mu[ev_type]
        contributions = []
        for pt, ptime in all_events:
            if ptime >= ev_time:
                break
            lam += alpha[ev_type, pt] * np.exp(-beta * (ev_time - ptime))

        if lam > 0:
            for pt, ptime in all_events:
                if ptime >= ev_time:
                    break
                contrib = np.exp(-beta * (ev_time - ptime)) / lam
                grad_alpha[ev_type, pt] += contrib

    return grad_mu, grad_alpha