"""Lipschitz-inflated split-conformal trajectory tube MVP."""
from dataclasses import dataclass
import numpy as np


def conformal_quantile(scores, alpha):
    """Finite-sample upper order statistic: rank ceil((m+1)(1-alpha))."""
    x = np.sort(np.asarray(scores, dtype=float))
    rank = int(np.ceil((len(x) + 1) * (1 - alpha)))
    return float(x[min(max(rank - 1, 0), len(x) - 1)])


def max_grid_slope(x, t):
    return float(np.max(np.abs(np.diff(x) / np.diff(t))))

@dataclass
class TrajectoryTube:
    q: float
    gamma: float

    @classmethod
    def calibrate(cls, trajectories, obs_times, alpha=0.1, Lhat=None,
                  predictor=lambda t, x: np.zeros_like(t), Lpred=0.0):
        scores = []
        for x, t in zip(trajectories, obs_times):
            scores.append(np.max(np.abs(x - predictor(t, x))))
        q = conformal_quantile(scores, alpha)
        if Lhat is None:
            raise ValueError("Lhat must be estimated on an independent high-frequency split")
        return cls(q=q, gamma=float(Lhat + Lpred))

    def radius(self, requested_times, observed_times):
        delta = np.min(np.abs(np.asarray(requested_times)[:, None] -
                              np.asarray(observed_times)[None, :]), axis=1)
        return self.q + self.gamma * delta, delta

    def contains(self, truth, center, requested_times, observed_times):
        r, _ = self.radius(requested_times, observed_times)
        return np.abs(np.asarray(truth) - np.asarray(center)) <= r + 1e-12