import torch


def trapezoidal_integral(y: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """
    Compute the trapezoidal integral of y with respect to t along the first dimension.

    y: (T, N) tensor of values to integrate
    t: (T,) tensor of time points corresponding to the first dimension of y

    Returns
    -------
    integral: (T, N) tensor of integrals for each voxel
    """
    integral = torch.zeros_like(y)
    integral[1:, ...] = torch.cumulative_trapezoid(y, t, dim=0)
    return integral


def quadratic_formula(
    A: torch.Tensor | float | int, B: torch.Tensor, C: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    discriminant = B**2 - 4 * A * C
    sqrt_discriminant = torch.sqrt(discriminant)
    x1 = (-B + sqrt_discriminant) / (2 * A)
    x2 = (-B - sqrt_discriminant) / (2 * A)
    return x1, x2


def interp1d_linear(
    x: torch.Tensor, y: torch.Tensor, x_new: torch.Tensor, *, eps: float
) -> torch.Tensor:
    if y.ndim not in {1, 2}:
        raise ValueError(f"y must be 1D or 2D. Got shape {tuple(y.shape)}.")

    x_new_clamped = x_new.clamp(min=x[0].item(), max=x[-1].item())

    idx = torch.searchsorted(x, x_new_clamped, right=True)
    idx = idx.clamp(1, x.numel() - 1)

    x0 = x[idx - 1]
    x1 = x[idx]
    alpha = (x_new_clamped - x0) / (x1 - x0).clamp_min(eps)

    if y.ndim == 1:
        y0 = y[idx - 1]
        y1 = y[idx]
        return y0 + alpha * (y1 - y0)

    y0 = y[idx - 1, :]
    y1 = y[idx, :]
    return y0 + alpha[:, None] * (y1 - y0)


def build_interpolation_grid(
    t_tac: torch.Tensor, t_aif: torch.Tensor, frame_duration_min: float
) -> torch.Tensor:
    if frame_duration_min <= 0.0:
        raise ValueError("frame_duration must be > 0.")

    t_start = torch.maximum(t_tac[0], t_aif[0])
    t_end = torch.minimum(t_tac[-1], t_aif[-1])
    if t_end <= t_start:
        raise ValueError("No temporal overlap between TAC and AIF time ranges.")

    step = t_tac.new_tensor(frame_duration_min)
    n_steps = int(torch.round((t_end - t_start) / step).item())
    n_points = max(2, n_steps + 1)
    return torch.linspace(
        t_start.item(), t_end.item(), n_points, device=t_tac.device, dtype=t_tac.dtype
    )
