from __future__ import annotations

import gc
from abc import ABC, abstractmethod
from contextlib import contextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Generator, Literal

import numpy as np
from funlib.geometry import Coordinate
from volara.utils import PydanticCoordinate, StrictBaseModel

if TYPE_CHECKING:
    import jax  # noqa: F401
    import torch


def _require_torch() -> None:
    try:
        import torch as _  # noqa: F811, F401
    except ImportError:
        raise ImportError(
            "PyTorch is required for this feature. "
            "Install it with: pip install volara-ml[torch]"
        ) from None


def _require_jax() -> None:
    try:
        import jax as _  # noqa: F811, F401
    except ImportError:
        raise ImportError(
            "JAX is required for this feature. "
            "Install it with: pip install volara-ml[jax]"
        ) from None


class Model(StrictBaseModel, ABC):
    """
    A base class for defining the common attributes and methods for all
    model types.
    """

    in_channels: int
    """
    The number of input channels to the model.
    """
    min_input_shape: PydanticCoordinate
    """
    The smallest possible input shape to the model. E.g. with UNets there
    is a minimum input size that satisfies the downsampling and all convolutional
    models have a minimum input size if using "valid" padding.
    """
    min_output_shape: PydanticCoordinate
    """
    The output shape to expect when providing the minimum input shape data.
    """
    min_step_shape: PydanticCoordinate
    """
    The minimum step size to increment the input shape by. This may be used
    to determine an optimal input/output shape for the model.
    """
    out_channels: list[int | None] | int | None
    """
    The number of channels in the output of the model. This can be an integer
    or list of integers if the model has multiple outputs. If `None` it is
    assumed that the number of output channels equals the number of input channels
    """
    out_range: tuple[float, float]
    """
    The range of the output data. This is used to convert between `np.uint8` and
    `np.float32` data types for efficient reading/writing of model outputs.
    """
    context_override: tuple[PydanticCoordinate, PydanticCoordinate] | None = None
    """
    An optional override for asymetrical context sizes. This lets you define
    a specific lower and upper context size for the model which must equal
    the expected context size of input_shape - output_shape
    """
    out_dtype: np.dtype = np.dtype(np.uint8)

    @property
    def context(self) -> Coordinate | tuple[Coordinate, Coordinate]:
        """
        The context required to make tile artifact free predictions
        with this model.
        """
        expected_context = self.min_input_shape - self.min_output_shape
        if self.context_override is not None:
            assert (
                self.context_override[0] + self.context_override[1]
            ) == expected_context, (
                f"Expected context override {self.context_override} to sum to {expected_context}, "
                f"but got {self.context_override[0] + self.context_override[1]}"
            )

        if self.context_override is not None:
            return self.context_override
        else:
            return expected_context // 2

    @abstractmethod
    def predict(self, device: Any) -> Generator[Callable[[np.ndarray], list[np.ndarray]], None, None]:
        """Context manager that yields a predict callable.

        Loads the model onto the given device on entry and frees GPU memory
        on exit (del model, clear caches, gc.collect).
        """
        pass

    @abstractmethod
    def select_device(self, worker_id: int) -> Any:
        """Return the device to use for the given worker."""
        pass

    @property
    @abstractmethod
    def eval_input_shape(self) -> Coordinate:
        """
        The input shape to use during prediction.
        """
        pass

    @property
    @abstractmethod
    def eval_output_shape(self) -> Coordinate:
        """
        The output shape to expect after providing data with shape
        `self.eval_input_shape()`
        """
        pass

    @property
    @abstractmethod
    def num_out_channels(self) -> list[int | None]:
        """
        The number of channels in each output prediction generated by this
        model. This is expected in a list to account for models that
        generate multiple output arrays.
        """
        pass

    def to_out_dtype(self, out_data: np.ndarray) -> np.ndarray:
        """
        A function defining the conversion function to go from model
        outputs to `np.uint8` for writing to zarr.
        """
        try:
            a_info = np.iinfo(self.out_dtype)
            a_min, a_max = a_info.min, a_info.max
        except ValueError:
            a_info = np.finfo(self.out_dtype)
            a_min, a_max = a_info.min, a_info.max
        return np.clip(
            ((out_data - self.out_range[0]) / (self.out_range[1] - self.out_range[0]))
            * a_max,
            a_min=a_min,
            a_max=a_max,
        ).astype(self.out_dtype)

    def from_uint8(self, data: np.ndarray) -> np.ndarray:
        """
        A function defining the conversion function to go from `np.uint8`
        back to `np.float32`
        """
        return (
            data.astype(np.float32) / 255 * (self.out_range[1] - self.out_range[0])
            + self.out_range[0]
        )


class TorchModel(Model):
    model_type: Literal["torch"] = "torch"

    def model_post_init(self, __context: object) -> None:
        _require_torch()

    save_path: Path
    """
    The path to a model saved using `torch.jit.save` or `torch.save`.
    """
    checkpoint_file: Path | None = None
    """
    An optional path to a checkpoint file containing the model weights.
    This can be provided either as a dictionary with key/value
    pair: "model_state_dict" and model.state_dict() or as a plain
    model.state_dict() dictionary."""
    pred_size_growth: PydanticCoordinate | None = None

    def model(self) -> "torch.nn.Module":
        import torch

        model = torch.load(self.save_path, map_location="cpu", weights_only=False)
        assert isinstance(model, torch.nn.Module), (
            f"Loading from {self.save_path} did not return a torch.nn.Module"
        )

        if self.checkpoint_file is not None:
            checkpoint = torch.load(
                self.checkpoint_file, map_location="cpu", weights_only=True
            )
            if "model_state_dict" in checkpoint:
                weights = checkpoint["model_state_dict"]
            else:
                weights = checkpoint
            model.load_state_dict(weights)

        return model

    @contextmanager
    def predict(self, device: Any) -> Generator[Callable[[np.ndarray], list[np.ndarray]], None, None]:
        import torch

        model = self.model()
        model.eval()
        model.to(device)

        def _predict(input_data: np.ndarray) -> list[np.ndarray]:
            with torch.no_grad():
                tensor = torch.as_tensor(input_data, device=device)
                out = model(tensor)  # noqa: F821
                if isinstance(out, tuple):
                    return [o.cpu().numpy() for o in out]
                return [out.cpu().numpy()]

        try:
            yield _predict
        finally:
            del model
            torch.cuda.empty_cache()
            gc.collect()

    def select_device(self, worker_id: int) -> Any:
        import torch

        n = torch.cuda.device_count()
        return "cpu" if n == 0 else f"cuda:{worker_id % n}"

    @property
    def eval_input_shape(self) -> Coordinate:
        input_shape = Coordinate(self.min_input_shape)

        if self.pred_size_growth is not None:
            assert np.sum(self.pred_size_growth % Coordinate(self.min_step_shape)) == 0
        if self.pred_size_growth is not None:
            input_shape = input_shape + self.pred_size_growth
        return input_shape

    @property
    def eval_output_shape(self) -> Coordinate:
        output_shape = Coordinate(self.min_output_shape)
        if self.pred_size_growth is not None:
            output_shape = output_shape + self.pred_size_growth
        return output_shape

    @property
    def num_out_channels(self) -> list[int | None]:
        num_channels = self.out_channels

        if isinstance(num_channels, int) or num_channels is None:
            return [num_channels]
        elif isinstance(num_channels, list):
            return num_channels


class JaxModel(Model):
    """A JAX/Flax model configuration.

    The model is loaded from a pickled Flax module (or apply function) and
    parameters are loaded from an orbax checkpoint directory or .msgpack file.
    All JAX imports are lazy so the module loads without JAX installed.
    """

    model_type: Literal["jax"] = "jax"
    model_path: Path
    """Path to a pickled Flax module (or apply function)."""
    params_path: Path | None = None
    """Path to an orbax checkpoint directory or .msgpack file with model parameters."""
    pred_size_growth: PydanticCoordinate | None = None

    def model_post_init(self, __context: object) -> None:
        _require_jax()

    def _load_model(self) -> Any:
        import pickle

        with open(self.model_path, "rb") as f:
            return pickle.load(f)

    def _load_params(self) -> Any:
        if self.params_path is None:
            return {}

        path = self.params_path
        if path.suffix == ".msgpack":
            from flax.serialization import from_bytes

            with open(path, "rb") as f:
                return from_bytes(None, f.read())
        else:
            import orbax.checkpoint as ocp

            checkpointer = ocp.StandardCheckpointer()
            return checkpointer.restore(path)

    @contextmanager
    def predict(self, device: Any) -> Generator[Callable[[np.ndarray], list[np.ndarray]], None, None]:
        import jax

        model = self._load_model()
        params = self._load_params()
        params = jax.device_put(params, device)

        @jax.jit
        def _apply(x: Any) -> Any:
            return model.apply(params, x)  # noqa: F821

        def _predict(input_data: np.ndarray) -> list[np.ndarray]:
            out = _apply(jax.numpy.asarray(input_data))
            if isinstance(out, tuple):
                return [np.asarray(o) for o in out]
            return [np.asarray(out)]

        try:
            yield _predict
        finally:
            del model, params
            jax.clear_caches()
            gc.collect()

    def select_device(self, worker_id: int) -> Any:
        import jax

        try:
            gpus = jax.devices("gpu")
        except RuntimeError:
            gpus = []
        if len(gpus) == 0:
            return jax.devices("cpu")[0]
        return gpus[worker_id % len(gpus)]

    @property
    def eval_input_shape(self) -> Coordinate:
        input_shape = Coordinate(self.min_input_shape)
        if self.pred_size_growth is not None:
            assert np.sum(self.pred_size_growth % Coordinate(self.min_step_shape)) == 0
            input_shape = input_shape + self.pred_size_growth
        return input_shape

    @property
    def eval_output_shape(self) -> Coordinate:
        output_shape = Coordinate(self.min_output_shape)
        if self.pred_size_growth is not None:
            output_shape = output_shape + self.pred_size_growth
        return output_shape

    @property
    def num_out_channels(self) -> list[int | None]:
        num_channels = self.out_channels
        if isinstance(num_channels, int) or num_channels is None:
            return [num_channels]
        elif isinstance(num_channels, list):
            return num_channels
