import numpy as np
import matplotlib.pyplot as plt
from radnn.plot.plot_base import PlotBase

class PlotImageGrid(PlotBase):
  def __init__(self, min=None, max=None, title=None):
    self.rows = []
    self.row_count = 0
    self.row_titles = []
    self.current_row = -1
    self.row_col_count = dict()
    self.max_col_count = 0
    self.min = min
    self.max = max
    self.title = title

  def add_row(self, row_title=None):
    self.current_row = self.row_count
    self.rows.append([])
    self.row_count = len(self.rows)
    self.row_titles.append(row_title)

  def add_column(self, image, image_title=None, color_map=None, aspect=None, extent=None):
    oRowColumns = self.rows[self.current_row]
    dImage = {"image": image, "title": image_title
      , "cmap": color_map, "aspect": aspect
      , "extend": extent}

    oRowColumns.append(dImage)
    self.rows[self.current_row] = oRowColumns
    nColCount = len(oRowColumns)

    self.row_col_count[self.current_row] = nColCount
    if nColCount > self.max_col_count:
      self.max_col_count = nColCount

  def prepare(self, images, grid=[4,4], figure_size=(15, 6), restrict_columns=None, color_map=None):
    nImageIndex = 0
    for nRows in range(grid[0]):
      self.add_row("")
      for nCols in range(grid[1]):
        nImage = images[nImageIndex, ...]
        self.add_column(nImage, "")
        nImageIndex += 1



    nColumns = restrict_columns
    if nColumns is None:
      nColumns = self.max_col_count
    fig, oSubplotGrid = plt.subplots(  nrows=self.row_count, ncols=nColumns
                                     , figsize=figure_size
                                     , subplot_kw={'xticks': [], 'yticks': []})
    bIsSingleRow = self.row_count == 1
    if bIsSingleRow:
      oSubplotGrid = oSubplotGrid[np.newaxis, ...]

    fig.suptitle(self.title)
    for nRowIndex, oRowColumns in enumerate(self.rows):
      if len(oRowColumns) > 0:
        sRowTitle = self.row_titles[nRowIndex]
        nRowImageCount = len(oRowColumns)
        #nIncr = nImageCount // nRowColumnCount
        nIncr = 1
        nImageIndex = 0
        for nColIndex in range(nColumns):
          bMustPlot = nColIndex < nRowImageCount
          #if (nIncr == 0) and (nColIndex > 0):
          #  bMustPlot = False

          if bMustPlot:
            dImage = oRowColumns[nImageIndex]
            oSubPlot = oSubplotGrid[nRowIndex, nColIndex]
            sTitle = dImage['title']
            if sTitle is not None:
              oSubPlot.title.set_text(sTitle)
            oSubPlot.set_xticks([])
            oSubPlot.set_yticks([])
            oSubPlot.imshow(dImage["image"], cmap=color_map,
                            aspect=dImage["aspect"], extent=dImage["extend"],
                            vmin=self.min, vmax=self.max
                            )

            if nColIndex == 0:
              if sRowTitle is not None:
                oSubPlot.text(0.0, 0.5, sRowTitle, transform=oSubPlot.transAxes,
                              horizontalalignment='right', verticalalignment='center',
                              fontsize=9, fontweight='bold')
          nImageIndex += nIncr
    fig.subplots_adjust(wspace=0.1, hspace=0.6)
    plt.tight_layout(pad=1.01)
    return self
