import matplotlib.pyplot as plt  # use the subpackage (a.k.a. namespace) with the alias "plt"
import numpy as np
from matplotlib import colors
from radnn.data import DataPreprocessingKind

# ====================================================================================================
class PlotDataset(object):  # class CPlot: object
  # --------------------------------------------------------------------------------------
  # Constructor
  def __init__(self, samples, labels
               , label_descriptions=["orange tree", "olive tree"]
               , colors=["darkorange", "darkseagreen"]
               , x_label="Feature 1"
               , y_label="Feature 2"
               ):
    # ................................................................
    # // Fields \\
    self.title = ""
    self.samples = samples
    self.labels = labels
    self.label_descriptions = label_descriptions
    self.colors = colors  # # https://matplotlib.org/3.1.0/gallery/color/named_colors.html
    self.x_label = x_label
    self.y_label = y_label
    # ................................................................

  # --------------------------------------------------------------------------------------
  def prepare(self, title, line_slope=None, line_intercept=None, limits_x=None, limits_y=None, preprocessed: DataPreprocessingKind | None = None):
    self.title = title

    # Two dimensional dataset for the scatter plot
    nXValues = self.samples[:, 0]
    nYValues = self.samples[:, 1]
    nLabels = self.labels

    oColorMap = colors.ListedColormap(self.colors)

    fig, ax = plt.subplots(figsize=(8, 8))
    plt.scatter(nXValues, nYValues, c=nLabels, cmap=oColorMap)

    plt.title(self.title)
    cb = plt.colorbar()
    nLoc = np.arange(0, max(nLabels), max(nLabels) / float(len(self.colors)))
    cb.set_ticks(nLoc)
    oLegend = [ f"{i}:{x}" for i,x in enumerate(self.label_descriptions)]
    cb.set_ticklabels(oLegend)

    if (line_slope is not None):
      x1 = np.min(nXValues)
      y1 = line_slope * x1 + line_intercept;
      x2 = np.max(nXValues)
      y2 = line_slope * x2 + line_intercept;
      oPlot1 = ax.plot([x1, x2], [y1, y2], 'r--', label="Decision line")
      oLegend = plt.legend(loc="upper left", shadow=True, fontsize='x-large')
      oLegend.get_frame().set_facecolor("lightyellow")

    if preprocessed == DataPreprocessingKind.MIN_MAX_NORMALIZE:
      ax.set_xlim((-0.05, 1.05))
      ax.set_ylim((-0.05, 1.05))
    if preprocessed == DataPreprocessingKind.STANDARDIZE:
      ax.set_xlim((-4.05, 4.05))
      ax.set_ylim((-4.05, 4.05))
    else:
      if limits_x is not None:
        ax.set_xlim(limits_x[0], limits_x[1])
      if limits_y is not None:
        ax.set_ylim(limits_y[0], limits_y[1])

    ax.set_xlabel(self.x_label)
    ax.set_ylabel(self.y_label)

    # plt.scatter(oDataset.Samples[:,0], oDataset.Samples[:,1])
    # , t, 'r--', t, t**2, 'bs', t, t**3, 'g^')
    return self
  # --------------------------------------------------------------------------------------
  def save(self, filename):
    plt.savefig(filename, bbox_inches='tight')
    return self
  # --------------------------------------------------------------------------------------
  def show(self):
    plt.show()
  # --------------------------------------------------------------------------------------


# ====================================================================================================

class CPlot(PlotDataset):
  # --------------------------------------------------------------------------------------
  def __init__(self, p_sTitle, p_nSamples, p_nLabels, p_sLabelDescriptions=["orange tree", "olive tree"]
               , p_sColors=["darkorange", "darkseagreen"]
               , p_sXLabel="Feature 1"
               , p_sYLabel="Feature 2"
               ):
    super().__init__(p_nSamples, p_nLabels, p_sLabelDescriptions, p_sColors, p_sXLabel, p_sYLabel)
    self.s_title = p_sTitle

  # --------------------------------------------------------------------------------------
  def Show(self, p_bIsMinMaxScaled=False, p_nLineSlope=None, p_nLineIntercept=None, p_nLimitsX=[-4, 4], p_nLimitsY=[-4, 4]):
    if p_bIsMinMaxScaled:
      self.prepare(self.s_title, p_nLineSlope, p_nLineIntercept, p_nLimitsX, p_nLimitsY, preprocessed=DataPreprocessingKind.MIN_MAX_NORMALIZE).show()
    else:
      self.prepare(self.s_title,  p_nLineSlope, p_nLineIntercept, p_nLimitsX, p_nLimitsY).show()
  # --------------------------------------------------------------------------------------
