# ======================================================================================
#
#     Rapid Deep Neural Networks
#
#     Licensed under the MIT License
# ______________________________________________________________________________________
# ......................................................................................

# Copyright (c) 2022-2026 Pantelis I. Kaplanoglou

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# .......................................................................................
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors

class PlotVisualizeDataset2D(object):
  # --------------------------------------------------------------------------------------------------------------------
  # Constructor
  def __init__(self, samples_2d, labels, title=None):
    # ................................................................
    # // Fields \\
    self.title = title
    self.samples = samples_2d
    self.labels = labels
    # ................................................................
  # --------------------------------------------------------------------------------------------------------------------
  def prepare(self, title=None, figure_size=(8, 8),
              is_min_max_scaled=False, line_slope=None, line_intercept=None,
              limits_x=[-4, 4], limits_y=[-4, 4]):
    if title is not None:
      self.title = title
      
    # Two dimensional dataset for the scatter plot
    nXValues = self.samples[:, 0]
    nYValues = self.samples[:, 1]
    nLabels = self.labels
    
    # https://matplotlib.org/3.1.0/gallery/color/named_colors.html
    oColors = ["darkorange", "darkseagreen"]
    oLabelDescriptions = ["orange tree", "olive tree"]
    oColorMap = colors.ListedColormap(oColors)
    
    fig, ax = plt.subplots(figsize=figure_size)
    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(oColors)))
    cb.set_ticks(nLoc)
    cb.set_ticklabels(oLabelDescriptions)
    
    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 is_min_max_scaled:
      ax.set_xlim((-0.05, 1.05))
      ax.set_ylim((-0.05, 1.05))
    else:
      ax.set_xlim(limits_x[0], limits_x[1])
      ax.set_ylim(limits_y[0], limits_y[1])
    
    ax.set_xlabel('Feature 1')
    ax.set_ylabel('Feature 2')
    
    return self
  
  # --------------------------------------------------------------------------------------------------------------------
  def save(self, filename):
    plt.savefig(filename, bbox_inches='tight')
    return self
  # --------------------------------------------------------------------------------------------------------------------
  def show(self):
    plt.show()
  # --------------------------------------------------------------------------------------------------------------------

