import numpy as np
import matplotlib.pyplot as plt
from .plot_base import PlotBase
from radnn.errors import ERR_PLOT_IMAGE_SHOULD_BE_NUMPY
from radnn.utils import capitalize


class PlotImage(PlotBase):
  # --------------------------------------------------------------------------------------
  def __init__(self, image, title=None):
    super().__init__()
    
    if isinstance(image, dict):
      self.image = image["sample"]
      nId = image["id"]
      nLabel = image["label"]
      sClassName = image["class_name"]
      if title is None:
        title = ""
        if nLabel is not None:
          title += f"[{nId}]"
        if sClassName is not None:
          title += f" {capitalize(sClassName)}"
    else:
      self.image = image
    
    assert isinstance(self.image, np.ndarray), ERR_PLOT_IMAGE_SHOULD_BE_NUMPY
    
    self.image = self.image.astype(np.uint8)
    
    self.color_map = None
    if self._is_grayscale():
      self.color_map = "gray"
    
    self.image = self.image.squeeze()
    self.title = title
  
  # --------------------------------------------------------------------------------------
  def _is_grayscale(self):
    nDims = list(self.image.shape)
    bIsGrayscale = True
    if len(nDims) == 4:
      if nDims[1] in [1, 3]:
        bIsGrayscale = nDims[0] == 1
        self.image = self.image.transpose(1, 2, 0)
      elif nDims[-1] in [1, 3]:
        bIsGrayscale = nDims[-1] == 1
    elif len(nDims) == 3:
      if nDims[0] in [1, 3]:
        bIsGrayscale = nDims[0] == 1
        self.image = self.image.transpose(1, 2, 0)
      elif nDims[-1] in [1, 3]:
        bIsGrayscale = nDims[-1] == 1
    return bIsGrayscale
  
  # --------------------------------------------------------------------------------------
  def prepare(self, color_map=None):
    if color_map is not None:
      self.color_map = color_map
    
    plt.clf()
    if self.title is not None:
      plt.suptitle(self.title)
    plt.imshow(self.image, cmap=self.color_map)
    return self
  # --------------------------------------------------------------------------------------


def plot_image_save(image, filename, fs, title=None):
  PlotImage(image, title).prepare().save(filename, fs=fs)
