from typing import Tuple

import numpy as np

from httomo.block_interfaces import BlockData, BlockTransfer, generic_array
from httomo.runner.auxiliary_data import AuxiliaryData
from httomo.utils import gpu_enabled, make_3d_shape_from_array, xp



[docs]
class BaseBlock(BlockData, BlockTransfer):
    """
    Base block class providing default implementations for the data
    transferring/getting/setting behaviour needed for a block type to be processed by
    implementors of `MethodWrapper`. Ie, this class provides default implementations for the
    `BlockTransfer` and `BlockData` protocols.

    Note that the data indexing behaviour described in `DataIndexing` is not implemented in
    this class. If the default implementations for data transferring/getting/setting in this
    class are acceptable: inherit from `BaseBlock`, override where necessary, and implement
    `DataIndexing` in order to implement the `Block` protocol.
    """

    def __init__(self, data: np.ndarray, aux_data: AuxiliaryData) -> None:
        self._data = data
        self._aux_data = aux_data

    def __dir__(self) -> list[str]:
        """Return only those properties that are relevant for the data"""
        return ["data", "angles", "angles_radians", "darks", "flats", "dark", "flat"]

    @property
    def data(self) -> generic_array:
        return self._data

    @data.setter
    def data(self, new_data: generic_array):
        self._data = new_data

    @property
    def aux_data(self) -> AuxiliaryData:
        return self._aux_data

    def _empty_aux_array(self):
        empty_shape = list(self._data.shape)
        return np.empty_like(self._data, shape=empty_shape)

    @property
    def angles(self) -> np.ndarray:
        return self._aux_data.get_angles()

    @angles.setter
    def angles(self, new_angles: np.ndarray):
        self._aux_data.set_angles(new_angles)

    @property
    def angles_radians(self) -> np.ndarray:
        return self.angles

    @angles_radians.setter
    def angles_radians(self, new_angles: np.ndarray):
        self.angles = new_angles

    @property
    def darks(self) -> generic_array:
        darks = self._aux_data.get_darks(self.is_gpu)
        if darks is None:
            darks = self._empty_aux_array()
        return darks

    @darks.setter
    def darks(self, darks: generic_array):
        self._aux_data.set_darks(darks)

    # alias
    @property
    def dark(self) -> generic_array:
        return self.darks

    @dark.setter
    def dark(self, darks: generic_array):
        self.darks = darks

    @property
    def flats(self) -> generic_array:
        flats = self._aux_data.get_flats(self.is_gpu)
        if flats is None:
            flats = self._empty_aux_array()
        return flats

    @flats.setter
    def flats(self, flats: generic_array):
        self._aux_data.set_flats(flats)

    # alias
    @property
    def flat(self) -> generic_array:
        return self.flats

    @flat.setter
    def flat(self, flats: generic_array):
        self.flats = flats

    @property
    def shape(self) -> Tuple[int, int, int]:
        """Shape of the data in this block"""
        return make_3d_shape_from_array(self._data)


[docs]
    def to_gpu(self):
        if not gpu_enabled:
            raise ValueError("no GPU available")
        self._data = xp.asarray(self.data, order="C")



[docs]
    def to_cpu(self):
        if not gpu_enabled:
            return
        self._data = xp.asnumpy(self.data, order="C")


    @property
    def is_gpu(self) -> bool:
        return not self.is_cpu

    @property
    def is_cpu(self) -> bool:
        return getattr(self._data, "device", None) is None