import warnings

import numpy as np
import torch

from .io_utils import (
    convert_time_to_minutes,
    infer_device,
    infer_dtype,
    normalize_tacs_shape,
    require_1d,
    to_tensor,
)
from .math_utils import build_interpolation_grid, interp1d_linear
from .solver import lls_3k, lls_3k_vB, lls_4k, lls_4k_vB
from .types import ArrayLike, DeviceLike, Scalar


def fit_compartment_model(
    tacs: ArrayLike,
    aif: ArrayLike,
    time: ArrayLike,
    *,
    model: str = "3k_vB",
    vB: Scalar | ArrayLike | None = None,
    aif_time: ArrayLike | None = None,
    interpolate: bool = True,
    frame_duration: float = 1.0 / 60.0,
    time_unit: str = "min",
    device: DeviceLike | None = None,
    dtype: torch.dtype | None = None,
    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[np.ndarray, np.ndarray, float]:
    """High-level interface for LLS compartment fitting.

    Accepts NumPy or torch inputs, runs fitting in torch, and always returns NumPy
    outputs as:
        (kinetic_params, fitted_tacs, mse)
    """

    resolved_device = infer_device(device, tacs, aif, time, aif_time, vB)
    resolved_dtype = infer_dtype(dtype, tacs, aif, time, aif_time)

    tacs_t = normalize_tacs_shape(
        to_tensor(tacs, device=resolved_device, dtype=resolved_dtype)
    )
    aif_t = to_tensor(aif, device=resolved_device, dtype=resolved_dtype)
    t_tac = to_tensor(time, device=resolved_device, dtype=resolved_dtype)
    t_aif = (
        t_tac
        if aif_time is None
        else to_tensor(aif_time, device=resolved_device, dtype=resolved_dtype)
    )

    if aif_t.ndim != 1:
        raise ValueError(f"aif must be 1D. Got shape {tuple(aif_t.shape)}.")
    if tacs_t.size(0) != t_tac.numel():
        raise ValueError("time and tacs must share the same number of samples.")
    if aif_t.numel() != t_aif.numel():
        raise ValueError("aif and aif_time must share the same number of samples.")

    require_1d("time", t_tac)
    require_1d("aif_time", t_aif)

    t_tac_min = convert_time_to_minutes(
        t_tac, time_unit=time_unit, warn_on_long_minutes=True
    )
    t_aif_min = convert_time_to_minutes(t_aif, time_unit=time_unit)

    if interpolate:
        t_model = build_interpolation_grid(t_tac_min, t_aif_min, frame_duration)
        tacs_model = interp1d_linear(t_tac_min, tacs_t, t_model, eps=eps)
        aif_model = interp1d_linear(t_aif_min, aif_t, t_model, eps=eps)
    else:
        if t_tac_min.shape != t_aif_min.shape or not torch.allclose(
            t_tac_min, t_aif_min, atol=eps, rtol=0.0
        ):
            raise ValueError("time and aif_time must match when interpolate=False.")
        t_model = t_tac_min
        tacs_model = tacs_t
        aif_model = aif_t

    model_key = model.lower().replace("-", "_")
    fixed_vb_models = {"3k", "lls_3k", "4k", "lls_4k"}
    fitted_vb_models = {"3k_vb", "lls_3k_vb", "4k_vb", "lls_4k_vb"}
    valid_models = fixed_vb_models | fitted_vb_models
    if model_key not in valid_models:
        raise ValueError(
            "Invalid model. Choose one of: "
            "'3k', '3k_vB', '4k', '4k_vB', 'lls_3k', 'lls_3k_vB', 'lls_4k', "
            "'lls_4k_vB'."
        )

    if model_key in fixed_vb_models:
        if vB is None:
            raise ValueError("vB must be provided for fixed-vB models ('3k' and '4k').")
        vb_t = to_tensor(vB, device=resolved_device, dtype=resolved_dtype)
    else:
        if vB is not None:
            warnings.warn(
                f"Ignoring vB for model='{model}'. This model estimates vB.",
                stacklevel=2,
            )
        vb_t = None

    kwargs = dict(
        reg=reg,
        eps=eps,
        column_scale=column_scale,
        nnls=nnls,
        nnls_max_iter=nnls_max_iter,
        nnls_tol=nnls_tol,
    )
    if model_key in {"3k", "lls_3k"}:
        kinetic_params, fitted_model, mse = lls_3k(
            tacs_model, aif_model, t_model, vB=vb_t, **kwargs
        )
    elif model_key in {"4k", "lls_4k"}:
        kinetic_params, fitted_model, mse = lls_4k(
            tacs_model, aif_model, t_model, vB=vb_t, **kwargs
        )
    elif model_key in {"3k_vb", "lls_3k_vb"}:
        kinetic_params, fitted_model, mse = lls_3k_vB(
            tacs_model, aif_model, t_model, **kwargs
        )
    else:
        kinetic_params, fitted_model, mse = lls_4k_vB(
            tacs_model, aif_model, t_model, **kwargs
        )

    if interpolate:
        fitted_tacs = interp1d_linear(t_model, fitted_model, t_tac_min, eps=eps)
    else:
        fitted_tacs = fitted_model

    kinetic_params_np = kinetic_params.detach().cpu().numpy()
    fitted_tacs_np = fitted_tacs.detach().cpu().numpy()
    mse_np = mse.detach().cpu().numpy()
    return kinetic_params_np, fitted_tacs_np, mse_np
