import warnings

import torch

from .models import (
    MODEL_3K,
    MODEL_4K,
    CompartmentModel,
    MODEL_3K_vB,
    MODEL_4K_vB,
)
from .types import DeviceLike, VBLike


def _validate_inputs(Cpet: torch.Tensor, Cp: torch.Tensor, t: torch.Tensor) -> None:
    assert Cpet.device == Cp.device == t.device, "All inputs must be on the same device"
    assert Cpet.dtype == Cp.dtype == t.dtype, "All inputs must have the same dtype"
    assert Cpet.ndim == 2, "Cpet must have shape (T, N)"
    assert Cp.ndim == 1, "Cp must have shape (T,)"
    assert t.ndim == 1, "t must have shape (T,)"
    assert Cpet.size(0) == Cp.size(0) == t.size(0), (
        "Mismatched time dimension in inputs"
    )


def _prepare_tissue_concentration(
    Cpet: torch.Tensor, Cp: torch.Tensor, vB: VBLike, eps: float
) -> tuple[torch.Tensor, torch.Tensor]:
    # Convert Cpet to Ct (tissue concentration) by removing blood contribution
    device = Cpet.device
    dtype = Cpet.dtype
    n_voxels = Cpet.size(1)

    vB_t = vB if torch.is_tensor(vB) else torch.tensor(vB, device=device, dtype=dtype)
    if vB_t.ndim == 0:  # Scalar
        vB_t = vB_t.clamp(0.0, 1.0 - eps)
        den = (1.0 - vB_t).clamp_min(eps)
        Ct = (Cpet - vB_t * Cp[:, None]) / den
    else:  # voxelwise (N,)
        vB_t = vB_t.to(device=device, dtype=dtype).clamp(0.0, 1.0 - eps)
        assert vB_t.ndim == 1 and vB_t.numel() == n_voxels, (
            "Vector vB must have shape (N,)"
        )
        den = (1.0 - vB_t).clamp_min(eps)  # (N,)
        Ct = (Cpet - Cp[:, None] * vB_t[None, :]) / den[None, :]

    return Ct, vB_t


def _solve_unconstrained(
    Gram: torch.Tensor, Y: torch.Tensor, *, device: DeviceLike
) -> torch.Tensor:
    try:
        return torch.linalg.solve(Gram, Y)
    except Exception as e:
        warnings.warn(f"Solve failed on {device} ({e}). Falling back to CPU float 64.")
        Gram_cpu = Gram.double().cpu()
        Y_cpu = Y.double().cpu()
        p_cpu = torch.linalg.solve(Gram_cpu, Y_cpu).float()
        return p_cpu.to(device)


def _solve_coefficients(
    X: torch.Tensor,
    Ct: torch.Tensor,
    *,
    reg: float,
    eps: float,
    column_scale: bool,
    nnls: bool,
    nnls_max_iter: int,
    nnls_tol: float,
    p_lower: torch.Tensor,
    p_upper: torch.Tensor,
) -> torch.Tensor:
    _, N, P = X.shape
    device = X.device
    dtype = X.dtype

    # Column scaling
    if column_scale:
        # per-voxel scaling for all columns (simple + robust)
        col_scale = torch.sqrt(torch.sum(X * X, dim=0)).clamp_min(eps)  # (N, P)
        Xs = X / col_scale[None, :, :]
    else:
        col_scale = None
        Xs = X

    # Gram matrix: X^T @ X
    Gram = torch.einsum("tni,tnj->nij", Xs, Xs)  # (N, P, P)

    # Apply Tikhonov regularization (ridge) to the Gram matrix
    Gram += reg * torch.eye(P, dtype=dtype, device=device).expand(N, P, P)

    # Right-hand side: X^T @ Ct
    Y = torch.einsum("tni,tn->ni", Xs, Ct)  # (N, P)

    # Non negative least squares with projected gradient descent
    if nnls:
        assert nnls_max_iter >= 1, "nnls_max_iter must be >= 1"
        assert nnls_tol > 0, "nnls_tol must be > 0"

        lower = p_lower.view(1, P).expand(N, P)
        upper = p_upper.view(1, P).expand(N, P)

        def project(p: torch.Tensor) -> torch.Tensor:
            return torch.minimum(torch.maximum(p, lower), upper)

        p = project(_solve_unconstrained(Gram, Y, device=device))

        # Step size from ||Gram||_inf upper bound of Lipschitz constant
        L_inf = (
            torch.abs(Gram).sum(dim=2).amax(dim=1, keepdim=True).clamp_min(eps)
        )  # (N,1)
        step = 1.0 / L_inf

        # Projected gradient iterations
        p_norm_floor = 1.0
        for _ in range(nnls_max_iter):
            grad = torch.bmm(Gram, p.unsqueeze(-1)).squeeze(-1) - Y  # (N, P)
            p_next = project(p - step * grad)

            # convergence
            rel = (p_next - p).norm(dim=1) / p.norm(dim=1).clamp_min(p_norm_floor)
            p = p_next
            if torch.max(rel) <= nnls_tol:
                break

        if not torch.isfinite(p).all():
            warnings.warn(
                "NNLS produced non-finite p values; consider increasing reg or enabling column_scale."
            )
    else:
        p = _solve_unconstrained(Gram, Y, device=device)

    # Unscale coefficients back if we scaled columns
    if column_scale:
        p = p / col_scale

    return p


def _lls_solve(
    Cpet: torch.Tensor,
    Cp: torch.Tensor,
    t: torch.Tensor,
    *,
    vB: VBLike | None,
    model: CompartmentModel,
    reg: float,
    eps: float,
    column_scale: bool,
    nnls: bool,
    nnls_max_iter: int,
    nnls_tol: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    _validate_inputs(Cpet, Cp, t)

    # Fixed-vB models solve in Ct-space, while fitted-vB models solve in Cpet-space.
    if model.fits_vB:
        y = Cpet
        X = model.build_design_matrix(Cpet, Cp, t)
        vB_t = None
    else:
        y, vB_t = _prepare_tissue_concentration(Cpet, Cp, vB, eps)
        X = model.build_design_matrix(y, Cp, t)

    p_lower, p_upper = model.coefficient_bounds(device=Cpet.device, dtype=Cpet.dtype)

    p = _solve_coefficients(
        X,
        y,
        reg=reg,
        eps=eps,
        column_scale=column_scale,
        nnls=nnls,
        nnls_max_iter=nnls_max_iter,
        nnls_tol=nnls_tol,
        p_lower=p_lower,
        p_upper=p_upper,
    )

    if model.fits_vB:
        # Compute fitted Cpet directly in Cpet-space
        fitted_Cpet = torch.einsum("tni,ni->tn", X, p)  # (T, N)
    else:
        # Compute fitted tissue concentration at each time point: X @ p
        fitted_Ct = torch.einsum("tni,ni->tn", X, p)  # (T, N)

        # Add back blood contribution to get fitted Cpet
        fitted_Cpet = fitted_Ct * (1.0 - vB_t) + Cp[:, None] * vB_t

    mse = torch.mean(torch.square(fitted_Cpet - Cpet), dim=0)  # (N,)

    kinetic_params = model.recover_kinetic_parameters(p, eps)
    return kinetic_params, fitted_Cpet, mse


def lls_3k(
    Cpet: torch.Tensor,
    Cp: torch.Tensor,
    t: torch.Tensor,
    *,
    vB: VBLike,
    reg: float = 1e-6,
    eps: float = 1e-8,
    column_scale: bool = True,
    nnls: bool = False,
    nnls_max_iter: int = 200,
    nnls_tol: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Linear least squares fitting of the irreversible two-tissue compartment model."""
    return _lls_solve(
        Cpet,
        Cp,
        t,
        vB=vB,
        model=MODEL_3K,
        reg=reg,
        eps=eps,
        column_scale=column_scale,
        nnls=nnls,
        nnls_max_iter=nnls_max_iter,
        nnls_tol=nnls_tol,
    )


def lls_4k(
    Cpet: torch.Tensor,
    Cp: torch.Tensor,
    t: torch.Tensor,
    *,
    vB: VBLike,
    reg: float = 1e-6,
    eps: float = 1e-8,
    column_scale: bool = True,
    nnls: bool = False,
    nnls_max_iter: int = 200,
    nnls_tol: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Linear least squares fitting of the reversible two-tissue compartment model."""
    return _lls_solve(
        Cpet,
        Cp,
        t,
        vB=vB,
        model=MODEL_4K,
        reg=reg,
        eps=eps,
        column_scale=column_scale,
        nnls=nnls,
        nnls_max_iter=nnls_max_iter,
        nnls_tol=nnls_tol,
    )


def lls_3k_vB(
    Cpet: torch.Tensor,
    Cp: torch.Tensor,
    t: torch.Tensor,
    *,
    reg: float = 1e-6,
    eps: float = 1e-8,
    column_scale: bool = True,
    nnls: bool = False,
    nnls_max_iter: int = 200,
    nnls_tol: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Linear least squares fitting of the irreversible model with fitted blood fraction."""
    return _lls_solve(
        Cpet,
        Cp,
        t,
        vB=None,
        model=MODEL_3K_vB,
        reg=reg,
        eps=eps,
        column_scale=column_scale,
        nnls=nnls,
        nnls_max_iter=nnls_max_iter,
        nnls_tol=nnls_tol,
    )


def lls_4k_vB(
    Cpet: torch.Tensor,
    Cp: torch.Tensor,
    t: torch.Tensor,
    *,
    reg: float = 1e-6,
    eps: float = 1e-8,
    column_scale: bool = True,
    nnls: bool = False,
    nnls_max_iter: int = 200,
    nnls_tol: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Linear least squares fitting of the reversible model with fitted blood fraction."""
    return _lls_solve(
        Cpet,
        Cp,
        t,
        vB=None,
        model=MODEL_4K_vB,
        reg=reg,
        eps=eps,
        column_scale=column_scale,
        nnls=nnls,
        nnls_max_iter=nnls_max_iter,
        nnls_tol=nnls_tol,
    )
