from abc import ABC, abstractmethod
from typing import ClassVar

import numpy as np
import torch

from .math_utils import quadratic_formula, trapezoidal_integral
from .types import DeviceLike


class ForwardCompartmentModel(ABC):
    """Forward model contract used for TAC simulation."""

    @abstractmethod
    def impulse_response(
        self, t: np.ndarray, K1: float, k2: float, k3: float = 0.0, k4: float = 0.0
    ) -> np.ndarray: ...

    def forward_cpet(
        self,
        Cp: np.ndarray,
        t: np.ndarray,
        K1: float,
        k2: float,
        vB: float,
        k3: float = 0.0,
        k4: float = 0.0,
    ) -> np.ndarray:
        cp = np.asarray(Cp)
        time = np.asarray(t)
        h = self.impulse_response(time, K1=K1, k2=k2, k3=k3, k4=k4)
        ct = np.convolve(h, cp, mode="full") * (time[1] - time[0])
        ct = ct[: len(time)]
        return (1 - vB) * ct + vB * cp


class CompartmentModel(ForwardCompartmentModel, ABC):
    """Model-specific hooks for design matrix construction and parameter recovery."""

    n_linear_params: ClassVar[int] = 0
    fits_vB: ClassVar[bool] = False

    @abstractmethod
    def build_design_matrix(
        self, Ct: torch.Tensor, Cp: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor: ...

    @abstractmethod
    def coefficient_bounds(
        self, *, device: DeviceLike | None, dtype: torch.dtype | str | None
    ) -> tuple[torch.Tensor, torch.Tensor]: ...

    @abstractmethod
    def recover_kinetic_parameters(
        self, p: torch.Tensor, eps: float
    ) -> torch.Tensor: ...


class Irreversible3KModel(CompartmentModel):
    n_linear_params: ClassVar[int] = 3

    def impulse_response(
        self, t: np.ndarray, K1: float, k2: float, k3: float = 0.0, k4: float = 0.0
    ) -> np.ndarray:
        del k4
        return _k3_impulse_response(t, K1=K1, k2=k2, k3=k3)

    def build_design_matrix(
        self, Ct: torch.Tensor, Cp: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        T, N = Ct.shape

        # Compute integrals for the design matrix
        I_Ct = trapezoidal_integral(Ct, t)  # (T, N)
        I_Cp = trapezoidal_integral(Cp, t)  # (T,)
        I2_Cp = trapezoidal_integral(I_Cp, t)  # (T,)

        # Design matrix: (T, N, 3) with columns [I_Ct, I_Cp, I2_Cp]
        X = torch.empty((T, N, 3), dtype=Ct.dtype, device=Ct.device)
        X[..., 0] = I_Ct
        X[..., 1] = I_Cp.unsqueeze(-1)
        X[..., 2] = I2_Cp.unsqueeze(-1)
        return X

    def coefficient_bounds(
        self, *, device: DeviceLike | None, dtype: torch.dtype | str | None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Bounds assume mapping:
        #   p1 <= 0,  p2 >= 0,  p3 >= 0
        p_lower = torch.tensor([-torch.inf, 0.0, 0.0], device=device, dtype=dtype)
        p_upper = torch.tensor([0.0, torch.inf, torch.inf], device=device, dtype=dtype)
        return p_lower, p_upper

    def recover_kinetic_parameters(self, p: torch.Tensor, eps: float) -> torch.Tensor:
        p1, p2, p3 = p[:, 0], p[:, 1], p[:, 2]

        # Recover kinetic parameters from linear coefficients
        K1 = p2
        k3 = p3 / (K1 + eps)
        k2 = -k3 - p1
        Ki = K1 * k3 / (k2 + k3 + eps)

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


class Irreversible3KModelvB(CompartmentModel):
    n_linear_params: ClassVar[int] = 4
    fits_vB: ClassVar[bool] = True

    def impulse_response(
        self, t: np.ndarray, K1: float, k2: float, k3: float = 0.0, k4: float = 0.0
    ) -> np.ndarray:
        del k4
        return _k3_impulse_response(t, K1=K1, k2=k2, k3=k3)

    def build_design_matrix(
        self, Ct: torch.Tensor, Cp: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        T, N = Ct.shape

        # Compute integrals for the design matrix
        I_Ct = trapezoidal_integral(Ct, t)  # (T, N)
        I_Cp = trapezoidal_integral(Cp, t)  # (T,)
        I2_Cp = trapezoidal_integral(I_Cp, t)  # (T,)

        # Design matrix: (T, N, 4) with columns [I_Ct, Cp, I_Cp, I2_Cp]
        X = torch.empty((T, N, 4), dtype=Ct.dtype, device=Ct.device)
        X[..., 0] = I_Ct
        X[..., 1] = Cp.unsqueeze(-1)
        X[..., 2] = I_Cp.unsqueeze(-1)
        X[..., 3] = I2_Cp.unsqueeze(-1)
        return X

    def coefficient_bounds(
        self, *, device: DeviceLike | None, dtype: torch.dtype | str | None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Bounds assume mapping:
        #   p1 <= 0, 0 <= p2 <= 1,  p3 >= 0,  p4 >= 0
        p_lower = torch.tensor([-torch.inf, 0.0, 0.0, 0.0], device=device, dtype=dtype)
        p_upper = torch.tensor(
            [0.0, 1.0, torch.inf, torch.inf], device=device, dtype=dtype
        )
        return p_lower, p_upper

    def recover_kinetic_parameters(self, p: torch.Tensor, eps: float) -> torch.Tensor:
        p1, p2, p3, p4 = p[:, 0], p[:, 1], p[:, 2], p[:, 3]

        # Recover kinetic parameters from linear coefficients
        K1 = (p3 + p1 * p2) / (1 - p2 + eps)
        k3 = p4 / ((1 - p2) * K1 + eps)
        k2 = -k3 - p1
        vB = p2
        Ki = K1 * k3 / (k2 + k3 + eps)

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


class Reversible4KModel(CompartmentModel):
    n_linear_params: ClassVar[int] = 4

    def impulse_response(
        self, t: np.ndarray, K1: float, k2: float, k3: float = 0.0, k4: float = 0.0
    ) -> np.ndarray:
        return _k4_impulse_response(t, K1=K1, k2=k2, k3=k3, k4=k4)

    def build_design_matrix(
        self, Ct: torch.Tensor, Cp: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        T, N = Ct.shape

        # Compute integrals for the design matrix
        I_Ct = trapezoidal_integral(Ct, t)  # (T, N)
        I2_Ct = trapezoidal_integral(I_Ct, t)  # (T, N)
        I_Cp = trapezoidal_integral(Cp, t)  # (T,)
        I2_Cp = trapezoidal_integral(I_Cp, t)  # (T,)

        # Design matrix: (T, N, 4) with columns [I_Ct, I2_Ct, I_Cp, I2_Cp]
        X = torch.empty((T, N, 4), dtype=Ct.dtype, device=Ct.device)
        X[..., 0] = I_Ct
        X[..., 1] = I2_Ct
        X[..., 2] = I_Cp.unsqueeze(-1)
        X[..., 3] = I2_Cp.unsqueeze(-1)
        return X

    def coefficient_bounds(
        self, *, device: DeviceLike | None, dtype: torch.dtype | str | None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Bounds assume mapping:
        #   p1 <= 0,  p2 <= 0,  p3 >= 0, p4 >= 0
        p_lower = torch.tensor(
            [-torch.inf, -torch.inf, 0.0, 0.0], device=device, dtype=dtype
        )
        p_upper = torch.tensor(
            [0.0, 0.0, torch.inf, torch.inf], device=device, dtype=dtype
        )
        return p_lower, p_upper

    def recover_kinetic_parameters(self, p: torch.Tensor, eps: float) -> torch.Tensor:
        p1, p2, p3, p4 = p[:, 0], p[:, 1], p[:, 2], p[:, 3]

        # Recover kinetic parameters from linear coefficients
        K1 = p3
        s = p4 / (p3 + eps)
        k2 = -p1 - s
        k4 = -p2 / (k2 + eps)
        k3 = s - k4
        Ki = K1 * k3 / (k2 + k3 + eps)

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


class Reversible4KModelvB(CompartmentModel):
    n_linear_params: ClassVar[int] = 5
    fits_vB: ClassVar[bool] = True

    def impulse_response(
        self, t: np.ndarray, K1: float, k2: float, k3: float = 0.0, k4: float = 0.0
    ) -> np.ndarray:
        return _k4_impulse_response(t, K1=K1, k2=k2, k3=k3, k4=k4)

    def build_design_matrix(
        self, Ct: torch.Tensor, Cp: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        T, N = Ct.shape

        # Compute integrals for the design matrix
        I_Ct = trapezoidal_integral(Ct, t)  # (T, N)
        I2_Ct = trapezoidal_integral(I_Ct, t)  # (T, N)
        I_Cp = trapezoidal_integral(Cp, t)  # (T,)
        I2_Cp = trapezoidal_integral(I_Cp, t)  # (T,)

        # Design matrix: (T, N, 5) with columns [I_Ct, I2_Ct, Cp, I_Cp, I2_Cp]
        X = torch.empty((T, N, 5), dtype=Ct.dtype, device=Ct.device)
        X[..., 0] = I_Ct
        X[..., 1] = I2_Ct
        X[..., 2] = Cp.unsqueeze(-1)
        X[..., 3] = I_Cp.unsqueeze(-1)
        X[..., 4] = I2_Cp.unsqueeze(-1)
        return X

    def coefficient_bounds(
        self, *, device: DeviceLike | None, dtype: torch.dtype | str | None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Bounds assume mapping:
        #   p1 <= 0,  p2 <= 0, 0 <= p3 <= 1, p4 >= 0, p5 >= 0
        p_lower = torch.tensor(
            [-torch.inf, -torch.inf, 0.0, 0.0, 0.0], device=device, dtype=dtype
        )
        p_upper = torch.tensor(
            [0.0, 0.0, 1.0, torch.inf, torch.inf], device=device, dtype=dtype
        )
        return p_lower, p_upper

    def recover_kinetic_parameters(self, p: torch.Tensor, eps: float) -> torch.Tensor:
        p1, p2, p3, p4, p5 = p[:, 0], p[:, 1], p[:, 2], p[:, 3], p[:, 4]

        # Recover kinetic parameters from linear coefficients
        K1 = (p4 + p1 * p3) / (1 - p3 + eps)
        k3 = p5 / (p4 + p1 * p3 + eps) + p3 * p2

        s1, s2 = quadratic_formula(1, p1 + k3, -p2)
        k2_mask = s1 > s2
        k2 = torch.where(k2_mask, s1, s2)
        k4 = torch.where(k2_mask, s2, s1)
        vB = p3

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

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


class OneTissue2KModel(ForwardCompartmentModel):
    def impulse_response(
        self, t: np.ndarray, K1: float, k2: float, k3: float = 0.0, k4: float = 0.0
    ) -> np.ndarray:
        del k3, k4
        return _k2_impulse_response(t, K1=K1, k2=k2)


MODEL_2K = OneTissue2KModel()
MODEL_3K = Irreversible3KModel()
MODEL_3K_vB = Irreversible3KModelvB()
MODEL_4K = Reversible4KModel()
MODEL_4K_vB = Reversible4KModelvB()


def _k2_impulse_response(t: np.ndarray, K1: float, k2: float) -> np.ndarray:
    return K1 * np.exp(-k2 * t)


def _k3_impulse_response(
    t: np.ndarray, K1: float, k2: float, k3: float
) -> np.ndarray:
    return K1 / (k2 + k3) * (k3 + k2 * np.exp(-(k2 + k3) * t))


def _k4_impulse_response(
    t: np.ndarray, K1: float, k2: float, k3: float, k4: float
) -> np.ndarray:
    disc = np.sqrt((k2 + k3 + k4) ** 2 - 4 * k2 * k4)
    a1 = 0.5 * (k2 + k3 + k4 + disc)
    a2 = 0.5 * (k2 + k3 + k4 - disc)

    b1 = (a1 - k3 - k4) * np.exp(-(a1 * t))
    b2 = (a2 - k3 - k4) * np.exp(-(a2 * t))

    return K1 / (a1 - a2) * (b1 - b2)


def make_compartment_model(num_k: int = 2):
    """Return forward compartment-model callable used for TAC simulation.

    The returned function maps (Cp, t, kinetic params) -> Cpet.
    """
    forward_models = {
        2: MODEL_2K,
        3: MODEL_3K,
        4: MODEL_4K,
    }
    model = forward_models.get(num_k)
    if model is None:
        def unsupported(*args, **kwargs):
            raise NotImplementedError(
                "Only compartment models with num_k in {2, 3, 4} are supported."
            )
        return unsupported
    return model.forward_cpet
