import numpy as np
import imageio.v3 as imageio
import matplotlib.pyplot as plt
from .core import naturalize_rgb_image
from .visualization import overlay_naturalness_heatmap


__all__ = ["GraNat"]


class GraNat:
    """
    A convenience wrapper around the GraNatPy naturalization pipeline.

    This class provides an interface to the core naturalization functions,
    plus loading, saving and image visualization without managing intermediate results manually.

    Typical usage::

        # From a file path
        result = GraNat.load_image("photo.tif").naturalize()
        # From an existing numpy array
        result = GraNat(img_array).naturalize()
        # Chained with visualization
        GraNat.load_image("photo.tif").naturalize().plot_heatmap()

    Attributes:
        image:   The input image as a uint8 numpy array of shape (H, W, 3)
                 or (H, W) for grayscale.
        result:  The naturalized image as a uint8 numpy array. None until
                 naturalize() has been called.
        nfs:     List of naturalization factors, one per channel. None until
                 naturalize() has been called.
        t1s:     List of T1 parameters, one per channel. None until
                 naturalize() has been called.
        t2s:     List of T2 parameters, one per channel. None until
                 naturalize() has been called.
        psnrs:   List of estimated PSNR values, one per channel. None until
                 naturalize() has been called.
    """

    def __init__(self, image: np.ndarray):

        self.image = image
        self.result = None
        self.nfs = None
        self.t1s = None
        self.t2s = None
        self.psnrs = None

    @classmethod
    def load_image(cls, path: str) -> "GraNat":
        """
        Load an image from disk and return a GraNat instance.
        Supports any format readable by imageio (TIFF, PNG, JPEG, etc.).

        Args:
            path: Path to the image file.

        Returns:
            A new GraNat instance initialised with the loaded image.

        Raises:
            FileNotFoundError: If no file exists at the given path.
            OSError: If the file exists but cannot be read as an image.

        Example::

            naturalizer = GraNat.load_image("photo.tif")
        """
        return cls(imageio.imread(path))

    def naturalize(self, show = False) -> "GraNat":
        """
        Run the naturalization pipeline on the loaded image.

        Processes each channel independently and stores the results in the
        instance attributes. Returns self to allow method chaining.

        Args:
            show: Whether to directly plot the image with matplotlib.
        Returns:
            self, to allow chaining e.g. GraNat.load_image(...).naturalize()
        Raises:
            ValueError: If the image array is not a valid 2D or 3D numpy array.
        Example::

            naturalizer = GraNat.load_image("photo.tif").naturalize()
            print(naturalizer.nfs)
        """
        self.result, self.nfs, self.t1s, self.t2s, self.psnrs = naturalize_rgb_image(self.image)
        if show:
            plt.imshow(self.result)
        return self

    def plot_heatmap(self, **kwargs):
        """
        Overlay a naturalness factor heatmap on the original image.

        Divides the image into a grid of cells and computes the average naturalization factor per cell,
        then renders a semi-transparent heatmap overlay. Keyword arguments are forwarded directly to
        overlay_naturalness_heatmap().

        Common keyword arguments:
            grid_rows (int):  Number of rows in the grid. Default 10.
            grid_cols (int):  Number of columns in the grid. Default 15.
            cmap (str):       Matplotlib colormap name. Default 'YlOrRd'.
            alpha (float):    Heatmap opacity in [0, 1]. Default 0.5.
            figsize (tuple):  Figure size as (width, height). Default (12, 8).

        Returns:
            A tuple of (fig, ax, heatmap_data) where heatmap_data is a 2D
            numpy array of shape (grid_rows, grid_cols) containing the
            average naturalization factor per cell.

        Raises:
            ValueError: If the image has not been loaded correctly.

        Example::

            fig, ax, heatmap = GraNat.load_image("photo.tif").plot_heatmap(grid_rows=20)
        """
        return overlay_naturalness_heatmap(self.image, **kwargs)

    def save(self, path: str) -> "GraNat":
        """
        Save the naturalized image to disk.

        Args:
            path: Output file path. The format is inferred from the file
                extension (e.g. '.png', '.tif', '.jpg').
        Returns:
            self, to allow chaining.
        Raises:
            RuntimeError: If naturalize() has not been called yet.
            OSError: If the file cannot be written to the given path.
        Example::

            GraNat.load_image("photo.tif").naturalize().save("output.png")
        """
        if self.result is None:
            raise RuntimeError("No result to save. Call naturalize() first.")
        imageio.imwrite(path, self.result)
        return self