"""Base class for dense data arrays with metadata."""
from __future__ import annotations

import abc
import gzip
import json
from pathlib import Path
from typing import Any, TypeVar

import numpy as np

from .chunk_types import FloatArray, PType
from .exceptions import ValidationError

T = TypeVar("T", bound="Chunk")


class Chunk(FloatArray):
    """
    Base class for dense data arrays with metadata.

    Parameters
    ----------
    data : FloatArray
        A numpy array of dense floats
    ptype : PType
        The parameterization of the dense data
    image_width : int, optional
        The width of the relevant image
    image_height : int, optional
        The height of the relevant image
    fps : float, optional
        The framerate of the chunk data (if tracking)
    object_ids:  list[str], optional
        Identifiers for the objects (if tracking)
    start_time: float, optional
        Start time of chunk with respect to video (if tracking)
    """

    empty_shape: tuple[int, ...]

    def __new__(
        cls: type[T],
        data: FloatArray,
        ptype: PType = PType.unknown,
        image_width: int | None = None,
        image_height: int | None = None,
        fps: float | None = None,
        object_ids: list[str] | None = None,
        start_time: float | None = None,
    ) -> T:
        """Instantiate a new chunk."""
        obj: T = np.asarray(data).view(cls)
        obj.ptype = ptype
        obj._image_width = image_width
        obj._image_height = image_height
        obj._fps = fps
        obj._object_ids = object_ids
        obj._start_time = start_time
        obj._scale = None
        return obj

    def __array_finalize__(self, obj: T | None) -> None:
        """Instantiate a new chunk created from a view."""
        if obj is None:
            return
        self.ptype: PType = getattr(obj, "ptype", PType.unknown)
        self._image_width: float | None = getattr(obj, "_image_width", None)
        self._image_height: float | None = getattr(obj, "_image_height", None)
        self._fps: float | None = getattr(obj, "_fps", None)
        self._object_ids: list[str] | None = getattr(obj, "_object_ids", None)
        self._start_time: float | None = getattr(obj, "_start_time", None)
        self._scale: FloatArray | None = getattr(obj, "_scale", None)
        self.validate()

    @abc.abstractmethod
    def validate(self) -> None:
        """Raise ValidationError for incorrect shape or values."""
        raise NotImplementedError

    @classmethod
    def empty(
        cls: type[T],
        ptype: PType = PType.unknown,
        image_width: int | None = None,
        image_height: int | None = None,
        fps: float | None = None,
        object_ids: list[str] | None = None,
        start_time: float | None = None,
    ) -> T:
        """Create an empty chunk."""
        data = np.zeros(cls.empty_shape) * np.nan
        return cls(
            data,
            ptype=ptype,
            image_width=image_width,
            image_height=image_height,
            fps=fps,
            object_ids=object_ids,
            start_time=start_time,
        )

    @property
    def array(self) -> FloatArray:
        """Dense data as an ndarray."""
        return self.view(np.ndarray)

    @property
    def image_width(self) -> float:
        """Image width."""
        if self._image_width is None:
            raise ValidationError("image_width not set")
        return self._image_width

    @property
    def image_height(self) -> float:
        """Image height."""
        if self._image_height is None:
            raise ValidationError("image_height not set")
        return self._image_height

    @property
    def fps(self) -> float:
        """Frames per second."""
        if self._fps is None:
            raise ValidationError("fps not set")
        return self._fps

    @property
    def object_ids(self) -> list[str]:
        """List of IDs for tracked objects."""
        if self._object_ids is None:
            return []
        return self._object_ids

    @property
    def start_time(self) -> float:
        """Start time of chunk with respect to a video."""
        if self._start_time is None:
            raise ValidationError("start_time not set")
        return self._start_time

    @property
    def duration(self) -> float:
        """Duration of the chunk in seconds."""
        raise NotImplementedError

    @property
    def metadata_kwargs(self) -> dict[str, Any]:
        """Metadata kwargs for downstream constructors."""
        return {
            "image_width": self._image_width,
            "image_height": self._image_height,
            "fps": self._fps,
            "object_ids": self._object_ids,
            "start_time": self._start_time,
        }

    def pad(self, shape: tuple[int, ...]) -> T:
        """Pad a chunk with NaN values."""
        try:
            assert len(shape) == self.ndim
            assert all(a <= b for a, b in zip(self.shape, shape))
        except AssertionError:
            raise ValidationError(
                f"Shape should be compatible with {self.shape} but got {shape}."
            )
        pad_width = tuple((0, b - a) for a, b in zip(self.shape, shape))
        data = np.pad(self.array, pad_width, "constant", constant_values=np.nan)
        return self.__class__(
            data,
            ptype=self.ptype,
            **self.metadata_kwargs,
        )

    def to_dict(self) -> dict[str, Any]:
        """Serialize chunk to a dict."""
        return {
            "data": np.where(np.isnan(self.array), np.array(None), self.array).tolist(),
            "classname": self.__class__.__name__,
            "ptype": self.ptype.name,
            **self.metadata_kwargs,
        }

    def to_file(self, path: str | Path) -> None:
        """Write serialized chunk to disk."""
        if not str(path).endswith(".json.gz"):
            raise ValidationError("Chunk file name must end with .json.gz")
        with gzip.open(path, "wt") as f:
            f.write(json.dumps(self.to_dict()))

    @classmethod
    def from_dict(
        cls: type[T], chunk_dict: dict[str, Any], dims: int | None = None
    ) -> T:
        """Create chunk from chunk dict."""
        data: FloatArray
        if np.array(chunk_dict["data"]).size == 0 and dims:
            if "tracking" in cls.__name__.lower():
                data = np.zeros((len(chunk_dict["data"]), 0, dims), "float64")
            else:
                data = np.zeros((len(chunk_dict["data"]), dims), "float64")
        else:
            data = np.array(chunk_dict["data"]).astype("float64")
        data[data == None] = np.nan
        return cls(
            data,
            ptype=PType(chunk_dict["ptype"]),
            image_width=chunk_dict["image_width"],
            image_height=chunk_dict["image_height"],
            fps=chunk_dict["fps"],
            object_ids=chunk_dict["object_ids"],
            start_time=chunk_dict.get("start_time"),  # Backwards compatibility
        )

    @classmethod
    def from_file(cls: type[T], path: str | Path) -> T:
        """Read chunk from disk."""
        with gzip.open(path, "rt") as f:
            return cls.from_dict(json.loads(f.read()))
