from abc import ABC, abstractmethod
from pathlib import Path
from typing import Literal

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


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.
    """

    @property
    def context(self) -> Coordinate:
        """
        The context required to make tile artifact free predictions
        with this model.
        """
        return (self.eval_input_shape - self.eval_output_shape) // 2

    @abstractmethod
    def model(self) -> torch.nn.Module:
        """
        A getter for a plain `torch.nn.Module` that can be called to
        generate the desired predictions. This should load the appropriate
        model weights.
        """
        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_uint8(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.
        """
        return np.clip(
            ((out_data - self.out_range[0]) / (self.out_range[1] - self.out_range[0]))
            * 255,
            0,
            255,
        ).astype(np.uint8)

    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"
    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

    @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
