import warnings

import numpy as np
from numpy.typing import NDArray

from ._torch_compat import ensure_torch, no_grad, torch_available

METHOD_PARAMS = {
    "patlak": ["Ki", "V0"],
    "lls3k": ["K1", "k2", "k3", "vB", "Ki"],
    "lls4k": ["K1", "k2", "k3", "k4", "vB", "Ki"],
}


def resolve_device(preference: str = "auto"):
    T = ensure_torch()
    if preference == "cpu":
        return T.device("cpu")
    if preference in ("auto", "mps") and T.backends.mps.is_available():
        return T.device("mps")
    if preference in ("auto", "cuda") and T.cuda.is_available():
        return T.device("cuda")
    return T.device("cpu")


def get_kinetic_maps(
    tacs: NDArray,
    aif: NDArray,
    time: NDArray,
    device: str = "auto",
    dtype: str = "float32",
    method: str = "patlak",
    n_batches: int = 1,
    *args,
    **kwargs,
):
    """
    Compute voxelwise kinetic model parameters from TACs and AIF using memory-lean
    linear least squares (LLS).

    Args
    ----
        tacs: (N, M) array of tissue TACs (time x voxels)
        aif : (N,)   array of arterial input function
        time: (N,)   array of time points (monotone increasing)
        device: 'cpu', 'cuda', 'mps', or 'auto' (default, prefers GPU if available)
        dtype: 'float32' (default), 'float64', or 'auto'.
               - 'float32': fastest / lowest memory
               - 'float64': closest numerical parity with legacy implementation
               - 'auto': use float64 if any input is float64, else float32
        method: 'patlak' (default), 'lls3k', or 'lls4k'
        n_batches: int >=1 or None; if >1, split voxels into n_batches chunks to reduce memory use
                     (may be slower). If None, process all voxels at once.
        *args, **kwargs: passed to the selected method

    Returns
    -------
        dict of parameter name to (M,) array, where M is the number of voxels.
        Parameter names depend on method; see METHOD_PARAMS.
    """
    T = ensure_torch()

    T.set_grad_enabled(False)
    device = resolve_device(device)

    if dtype == "float32":
        target_dtype = T.float32
    elif dtype == "float64":
        target_dtype = T.float64
    elif dtype == "auto":
        in_dtypes = (np.asarray(tacs).dtype, np.asarray(aif).dtype, np.asarray(time).dtype)
        target_dtype = (
            T.float64 if any(np.issubdtype(dt, np.float64) for dt in in_dtypes) else T.float32
        )
    else:
        raise ValueError("dtype must be one of {'float32', 'float64', 'auto'}")

    tacs = T.from_numpy(tacs).to(target_dtype, copy=False).to(device)
    aif = T.from_numpy(aif).to(target_dtype, copy=False).to(device)
    t = T.from_numpy(time).to(target_dtype, copy=False).to(device)

    if method == "lls3k":
        base_func = lls3k_streaming
    elif method == "lls4k":
        base_func = lls4k_streaming
    elif method == "patlak":
        base_func = patlak_streaming
    else:
        raise ValueError(f"Unknown method: {method}")

    def run(Ct):
        with no_grad():
            return base_func(Ct, aif, t, *args, **kwargs)

    if n_batches is None or n_batches < 2:
        params = run(tacs)  # (N, M) expected
    else:
        chunks = T.tensor_split(tacs, n_batches, dim=1)
        outs = []
        for c in chunks:
            out = run(c)  # (P, m_batch)
            outs.append(out.to("cpu"))  # move off-GPU immediately if using MPS
            if device == "mps":
                T.mps.empty_cache()
        params = T.cat(outs, dim=1)

    param_names = METHOD_PARAMS[method]
    return {name: params[i].cpu().numpy() for i, name in enumerate(param_names)}


@no_grad()
def lls3k_streaming(Ct, Cp, t, reg=1e-6, cpu_solve_fallback=True):
    """
    Memory-lean LLS for the irreversible 3K (two-tissue) model.

    Ct: (N, M) tissue TACs           (time x voxels)
    Cp: (N,)   arterial input
    t : (N,)   time (monotone)
    Returns: params (5, M) in order [K1, k2, k3, vB, Ki]
    """
    T = ensure_torch()

    assert Ct.ndim == 2
    N, M = Ct.shape
    device, dtype = Ct.device, Ct.dtype

    # Accumulators
    Gram = T.zeros((M, 4, 4), dtype=dtype, device=device)
    Y = T.zeros((M, 4), dtype=dtype, device=device)

    # Running trapezoid integrals
    I_Cp = T.zeros((), dtype=dtype, device=device)  # scalar
    I2_Cp = T.zeros((), dtype=dtype, device=device)  # scalar
    I_Ct = T.zeros((M,), dtype=dtype, device=device)  # per-voxel

    # Loop: update integrals to time n, then accumulate LS terms at time n
    for n in range(1, N):
        dt = t[n] - t[n - 1]
        # update integrals (trapezoid)
        I_Cp = I_Cp + 0.5 * (Cp[n] + Cp[n - 1]) * dt
        I2_Cp = I2_Cp + I_Cp * dt
        I_Ct = I_Ct + 0.5 * (Ct[n, :] + Ct[n - 1, :]) * dt

        # features at time n
        f1 = I_Ct  # (M,)
        f2 = Cp[n]  # scalar
        f3 = I_Cp  # scalar
        f4 = I2_Cp  # scalar

        # accumulate Gram = sum x x^T  (exploit broadcasting of scalars)
        Gram[:, 0, 0] += f1 * f1
        Gram[:, 0, 1] += f1 * f2
        Gram[:, 1, 0] += f1 * f2
        Gram[:, 0, 2] += f1 * f3
        Gram[:, 2, 0] += f1 * f3
        Gram[:, 0, 3] += f1 * f4
        Gram[:, 3, 0] += f1 * f4

        Gram[:, 1, 1] += f2 * f2
        Gram[:, 1, 2] += f2 * f3
        Gram[:, 2, 1] += f2 * f3
        Gram[:, 1, 3] += f2 * f4
        Gram[:, 3, 1] += f2 * f4

        Gram[:, 2, 2] += f3 * f3
        Gram[:, 2, 3] += f3 * f4
        Gram[:, 3, 2] += f3 * f4

        Gram[:, 3, 3] += f4 * f4

        # RHS: Y = sum x * b  with b = Ct[n]
        b = Ct[n, :]  # (M,)
        Y[:, 0] += f1 * b
        Y[:, 1] += f2 * b
        Y[:, 2] += f3 * b
        Y[:, 3] += f4 * b

    # Regularize
    eye4 = T.eye(4, dtype=dtype, device=device).expand(M, 4, 4)
    G = Gram + reg * eye4

    # Prefer Cholesky (SPD) and add tiny jitter if needed
    try:
        L, info = T.linalg.cholesky_ex(G)
        if T.any(info != 0):
            bad = info != 0
            # add extra jitter just for the failing batch elements
            G[bad] = G[bad] + (10 * reg) * T.eye(4, dtype=dtype, device=device)
            L = T.linalg.cholesky(G)
        p_all = T.cholesky_solve(Y.unsqueeze(-1), L).squeeze(-1)  # (M, 4)

    except Exception as e:
        warnings.warn(
            f"Cholesky failed on {device} ({e}). Falling back to CPU float 64."
        )
        if cpu_solve_fallback:
            G_cpu = G.double().cpu()
            Y_cpu = Y.double().cpu()
            p_cpu = T.linalg.solve(G_cpu, Y_cpu.unsqueeze(-1)).squeeze(-1).float()
            p_all = p_cpu.to(device)
        else:
            raise

    p1, p2, p3, p4 = p_all[:, 0], p_all[:, 1], p_all[:, 2], p_all[:, 3]

    denom = p1 * p2 + p3 + 1e-8
    K1 = (p1 * p2 + p3) / (1 - p2 + 1e-8)
    k2 = -p1 - p4 / denom
    k3 = p4 / denom
    vB = p2
    Ki = K1 * k3 / (k2 + k3 + 1e-8)

    return T.stack([K1, k2, k3, vB, Ki], dim=0)


@no_grad()
def lls4k_streaming(
    Ct,
    Cp,
    t,
    reg=1e-6,
    cpu_solve_fallback=True,
    eps=1e-8,
    solver="cholesky",
    nnls=False,
    nnls_max_iter=200,
    nnls_tol=1e-6,
):
    """
    Memory-lean LLS for the reversible 4-K model.

    Args
    ----
      Ct: (N, M) tissue TACs
      Cp: (N,)    arterial input function
      t : (N,)    time (monotone)
      reg:        Tikhonov ridge on Gram diagonals
      cpu_solve_fallback: if True, fall back to CPU float64 solve when MPS cholesky fails
      solver:     'cholesky' (default) or 'solve'. Use 'solve' for closest legacy parity.
      nnls:       if True, solve ridge-regularized sign-constrained least squares
                  on streamed normal equations (legacy arg name kept for compatibility):
                  p1<=0, p2<=0, p3 in [0,1], p4>=0, p5>=0.
      nnls_max_iter: max iterations for projected-gradient constrained solve
      nnls_tol:      relative update tolerance for projected-gradient convergence

    Returns
    -------
      params: (6, M) [K1, k2, k3, k4, vB, Ki]
    """
    T = ensure_torch()

    assert Ct.ndim == 2
    N, M = Ct.shape
    device, dtype = Ct.device, Ct.dtype

    # Accumulators
    Gram = T.zeros((M, 5, 5), dtype=dtype, device=device)
    Y = T.zeros((M, 5), dtype=dtype, device=device)

    # Running trapezoids
    I_Cp = T.zeros((), dtype=dtype, device=device)  # scalar
    I2_Cp = T.zeros((), dtype=dtype, device=device)  # scalar
    I_Ct = T.zeros((M,), dtype=dtype, device=device)  # per-voxel
    I2_Ct = T.zeros((M,), dtype=dtype, device=device)  # per-voxel

    def _accumulate_row(f1, f2, f3, f4, f5, b):
        # vector-vector
        Gram[:, 0, 0] += f1 * f1
        Gram[:, 0, 1] += f1 * f2
        Gram[:, 1, 0] += f1 * f2
        Gram[:, 1, 1] += f2 * f2

        # vector-scalar
        Gram[:, 0, 2] += f1 * f3
        Gram[:, 2, 0] += f1 * f3
        Gram[:, 0, 3] += f1 * f4
        Gram[:, 3, 0] += f1 * f4
        Gram[:, 0, 4] += f1 * f5
        Gram[:, 4, 0] += f1 * f5

        Gram[:, 1, 2] += f2 * f3
        Gram[:, 2, 1] += f2 * f3
        Gram[:, 1, 3] += f2 * f4
        Gram[:, 3, 1] += f2 * f4
        Gram[:, 1, 4] += f2 * f5
        Gram[:, 4, 1] += f2 * f5

        # scalar-scalar (same for all voxels)
        s33 = f3 * f3
        s34 = f3 * f4
        s35 = f3 * f5
        s44 = f4 * f4
        s45 = f4 * f5
        s55 = f5 * f5
        Gram[:, 2, 2] += s33
        Gram[:, 2, 3] += s34
        Gram[:, 3, 2] += s34
        Gram[:, 2, 4] += s35
        Gram[:, 4, 2] += s35
        Gram[:, 3, 3] += s44
        Gram[:, 3, 4] += s45
        Gram[:, 4, 3] += s45
        Gram[:, 4, 4] += s55

        # RHS: Y = sum x * b
        Y[:, 0] += f1 * b
        Y[:, 1] += f2 * b
        Y[:, 2] += f3 * b
        Y[:, 3] += f4 * b
        Y[:, 4] += f5 * b

    # Include the n=0 row so streaming matches the explicit design-matrix method.
    _accumulate_row(I_Ct, I2_Ct, Cp[0], I_Cp, I2_Cp, Ct[0, :])

    for n in range(1, N):
        dt = t[n] - t[n - 1]

        # Integrals: trapezoid for first and second integrals (matches double-cumtrapz).
        I_Cp_prev = I_Cp
        I_Ct_prev = I_Ct
        I_Cp = I_Cp + 0.5 * (Cp[n] + Cp[n - 1]) * dt
        Ct_mid = 0.5 * (Ct[n, :] + Ct[n - 1, :])  # (M,)
        I_Ct = I_Ct + Ct_mid * dt
        I2_Cp = I2_Cp + 0.5 * (I_Cp_prev + I_Cp) * dt
        I2_Ct = I2_Ct + 0.5 * (I_Ct_prev + I_Ct) * dt

        # Features at time n
        f1 = I_Ct  # (M,)
        f2 = I2_Ct  # (M,)
        f3 = Cp[n]  # scalar
        f4 = I_Cp  # scalar
        f5 = I2_Cp  # scalar

        _accumulate_row(f1, f2, f3, f4, f5, Ct[n, :])

    # Regularize + solve (prefer Cholesky)
    eye5 = T.eye(5, dtype=dtype, device=device).expand(M, 5, 5)
    G = Gram + reg * eye5

    if solver not in {"cholesky", "solve"}:
        raise ValueError("solver must be one of {'cholesky', 'solve'}")
    if nnls_max_iter < 1:
        raise ValueError("nnls_max_iter must be >= 1")
    if nnls_tol <= 0:
        raise ValueError("nnls_tol must be > 0")

    def _solve_unconstrained():
        if solver == "solve":
            try:
                return T.linalg.solve(G, Y.unsqueeze(-1)).squeeze(-1)  # (M, 5)
            except Exception as e:
                warnings.warn(
                    f"Solve failed on {device} ({e}). Falling back to CPU float 64."
                )
                if cpu_solve_fallback:
                    return (
                        T.linalg.solve(G.double().cpu(), Y.double().cpu().unsqueeze(-1))
                        .squeeze(-1)
                        .to(dtype)
                        .to(device)
                    )
                raise

        try:
            L, info = T.linalg.cholesky_ex(G)
            if T.any(info != 0):
                bad = info != 0
                G[bad] = G[bad] + (10 * reg) * T.eye(5, dtype=dtype, device=device)
                L = T.linalg.cholesky(G)
            return T.cholesky_solve(Y.unsqueeze(-1), L).squeeze(-1)  # (M, 5)
        except Exception as e:
            warnings.warn(
                f"Cholesky failed on {device} ({e}). Falling back to CPU float 64."
            )
            if cpu_solve_fallback:
                return (
                    T.linalg.solve(G.double().cpu(), Y.double().cpu().unsqueeze(-1))
                    .squeeze(-1)
                    .to(dtype)
                    .to(device)
                )
            raise

    if nnls:
        # Sign-consistent coefficient bounds for the 4K linearization:
        # p1<=0, p2<=0, p3 in [0,1], p4>=0, p5>=0.
        lower = T.tensor(
            [-T.inf, -T.inf, 0.0, 0.0, 0.0], dtype=dtype, device=device
        ).expand(M, 5)
        upper = T.tensor(
            [0.0, 0.0, 1.0 - eps, T.inf, T.inf], dtype=dtype, device=device
        ).expand(M, 5)

        def _project(x):
            return T.minimum(T.maximum(x, lower), upper)

        # Initialize with unconstrained ridge solution projected to feasible set.
        p_all = _project(_solve_unconstrained())
        # Conservative per-voxel step from ||G||_inf upper bound of Lipschitz constant.
        L_inf = T.abs(G).sum(dim=2).amax(dim=1, keepdim=True).clamp_min(eps)
        step = 1.0 / L_inf

        for _ in range(nnls_max_iter):
            grad = T.bmm(G, p_all.unsqueeze(-1)).squeeze(-1) - Y
            p_next = _project(p_all - step * grad)

            rel = (p_next - p_all).norm(dim=1) / p_all.norm(dim=1).clamp_min(1.0)
            p_all = p_next
            if T.max(rel) <= nnls_tol:
                break
    else:
        p_all = _solve_unconstrained()

    p1, p2, p3, p4, p5 = (p_all[:, i] for i in range(5))

    # Conversions
    K1 = (p4 + p1 * p3) / (1 - p3 + eps)
    k3 = p5 / (p4 + p1 * p3 + eps) - p3 * p2

    # quadratic for k2,k4
    B = p1 + k3
    C = -p2
    disc = T.clamp(B * B - 4 * C, min=0.0)
    sqrt_disc = T.sqrt(disc)
    s1 = 0.5 * (-B + sqrt_disc)
    s2 = 0.5 * (-B - sqrt_disc)
    k2_mask = s1 > s2
    k2 = T.where(k2_mask, s1, s2)
    k4 = T.where(k2_mask, s2, s1)

    vB = p3

    Ki = K1 * k3 / (k2 + k3 + eps)

    return T.stack([K1, k2, k3, k4, vB, Ki], dim=0)


@no_grad()
def patlak_streaming(Ct, Cp, t, break_index=None, break_time=None, eps=1e-8):
    """
    Memory-lean Patlak regression per voxel.

    Args
    ----
      Ct: (N, M)
      Cp: (N,)
      t : (N,)
      break_index: optional int index to start Patlak linear phase
      break_time : optional float (same units as t); use first n with t[n] >= break_time
      eps: Cp threshold for valid frames

    Returns
    -------
      params: (2, M) [Ki (slope), V0 (intercept)]
    """
    T = ensure_torch()

    assert Ct.ndim == 2
    N, M = Ct.shape
    device, dtype = Ct.device, Ct.dtype

    # Determine start index
    if break_index is not None:
        n0 = int(break_index)
    elif break_time is not None:
        # first index with t[n] >= break_time
        n0 = (
            int(T.nonzero(t >= break_time, as_tuple=False)[0].item())
            if T.any(t >= break_time)
            else N
        )  # (=> no frames)
    else:
        n0 = 0
    n0 = max(0, min(n0, N))  # clamp

    # Running integral of Cp
    I_Cp = T.zeros((), dtype=dtype, device=device)

    # OLS accumulators (Sx, Sx2, N) are scalars; Sy and Sxy are per-voxel
    Sx = T.zeros((), dtype=dtype, device=device)  # sum X
    Sx2 = T.zeros((), dtype=dtype, device=device)  # sum X^2
    Sy = T.zeros((M,), dtype=dtype, device=device)  # sum Y
    Sxy = T.zeros((M,), dtype=dtype, device=device)  # sum X*Y
    N_eff = T.zeros((), dtype=dtype, device=device)  # count

    for n in range(1, N):
        # update I_Cp up to frame n (trapezoid)
        dt = t[n] - t[n - 1]
        I_Cp = I_Cp + 0.5 * (Cp[n] + Cp[n - 1]) * dt

        if n < n0:
            continue

        cpn = Cp[n]
        if cpn <= eps:
            continue  # skip invalid frame

        Xn = I_Cp / cpn  # scalar
        Yn = Ct[n, :] / cpn  # (M,)

        Sx = Sx + Xn
        Sx2 = Sx2 + Xn * Xn
        Sy = Sy + Yn
        Sxy = Sxy + Xn * Yn
        N_eff = N_eff + 1

    # Handle the no-data case
    if N_eff.item() == 0:
        Ki = T.zeros((M,), dtype=dtype, device=device)
        V0 = T.zeros((M,), dtype=dtype, device=device)
        return T.stack([Ki, V0], dim=0)

    den = N_eff * Sx2 - Sx * Sx
    den = den + eps  # ridge to avoid division by zero

    Ki = (N_eff * Sxy - Sx * Sy) / den
    V0 = (Sy - Ki * Sx) / (N_eff + eps)

    return T.stack([Ki, V0], dim=0)
