from sklearn.cluster import KMeans
  import pickle
  import numpy as np
  import matplotlib.pyplot as plt
  from sklearn.manifold import TSNE
  from abstract_clustering import AbstractClustering
  
  class kmeansMahalanobis():

      def __init__(self, constrained: bool = False):

          """
  
          """
          self.C = None
          self.L = None
          self.K = None

          self.constrained = constrained

  
      def predict(self, features):
          """
  
          @param features:
          @return:
          """
          N = features.shape[0]
          distances = np.zeros((N, self.K))
          for n in range(N):
              for k in range(self.K):
                  distances[n][k] = self._dist(features[n], self.C[k], self.L[k])

          closest_cluster = np.argmin(distances, axis=1)
          return closest_cluster
  
      def load(self, model_path):
          """
  
          @param model_path:
          @return:
          """
          data = None

          with open(model_path, "rb") as f:
              data = pickle.load(f)

          if data is None:
              raise Exception("Le modèle n'a pas pu être chargé")
          else:
              self.C = data["C"]
              self.L = data["L"]
              self.K = data["K"]

              self.constrained = data["constrained"]

  
      def save(self, modelpath: str):
          """
  
          @param modelpath:
          @return:
          """
          data = {
              "C": self.C,
              "L": self.L,

              "K": self.K,
              "constrained": self.constrained

          }
          with open(modelpath, "wb") as f:
              pickle.dump(data, f)

      def fit(self, features, k: int, tol: float, ninit: int, maxiter: int=300, debug: bool=False):
          results = []
          for i in range(ninit):
              results.append(self._train(features, k, tol, maxiter, debug))
          losses = [v["loss"] for v in results]
          best = results[losses.index(min(losses))]
          if debug:
              print(f"best: {best['loss']} loss")
          self.C = best["C"]
          self.L = best["L"]
          self.K = best["K"]

  
      def _initialize_model(self, X, number_clusters):
          d = X.shape[1]
          C = X[np.random.choice(X.shape[0], number_clusters)]
          L = np.zeros((number_clusters, d, d))
          for k in range(number_clusters):
              L[k] = np.identity(d)
          return C, L
  
      def _dist(self, a, b, l):
          '''

          Distance euclidienne with mahalanobis

          '''
          a = np.reshape(a, (-1, 1))
          b = np.reshape(b, (-1, 1))

          result = np.transpose(a - b).dot(l).dot(a - b)[0][0]

          return result
  
      def _plot_iteration(self, iteration, points, clusters, centers):
          fig = plt.figure()
          ax = fig.add_subplot(111)
          scatter = ax.scatter(points[:, 0], points[:, 1], c=clusters, s=50)

          ax.scatter(centers[:, 0], centers[:, 1], s=50, c='red', marker='+')

          ax.set_xlabel('x')
          ax.set_ylabel('y')
          plt.colorbar(scatter)

          plt.savefig("test_" + str(iteration) + ".pdf")

      def _train(self, features, K: int, tol: float, maxiter: int, debug: bool=False):

          X = features
          N = X.shape[0]
          d = X.shape[1]
  
          C, L = self._initialize_model(X, K)
          self.C = C
          self.L = L
          self.K = K
  
          end_algo = False
          i = 0
          while not end_algo:

              if debug:
                  print("Iteration: ", i)

              # Calcul matrix distance

              distances = np.zeros((N, self.K))

  
              for n in range(N):
                  for k in range(self.K):
                      distances[n][k] = self._dist(X[n], self.C[k], self.L[k])

              closest_cluster = np.argmin(distances, axis=1)

              loss = np.sum(distances[np.arange(len(distances)), closest_cluster])
              if debug:
                  print(f"loss {loss}")

  
              # -- Debug tool ----------------------

              if debug and i % 1 == 0:

                  # TSNE if needed
                  X_embedded = np.concatenate((X, self.C), axis=0)
                  if d > 2:

                      X_embedded = TSNE(n_components=2).fit_transform(np.concatenate((X, self.C), axis=0))

                  # Then plot
                  self._plot_iteration(
                      i,
                      X_embedded[:X.shape[0]],
                      closest_cluster,
                      X_embedded[X.shape[0]:]
                  )

              # ------------------------------------

              old_c = self.C.copy()

              for k in range(self.K):

                  # Find subset of X with values closed to the centroid c_k.
                  X_sub = np.where(closest_cluster == k)
                  X_sub = np.take(X, X_sub[0], axis=0)

                  if X_sub.shape[0] == 0:
                      continue

                  C_new = np.mean(X_sub, axis=0)
  
                  # -- COMPUTE NEW LAMBDA (here named K) --

                  K_new = np.zeros((self.L.shape[1], self.L.shape[2]))
                  tmp = np.zeros((self.L.shape[1], self.L.shape[2]))

                  for x in X_sub:
                      x = np.reshape(x, (-1, 1))
                      c_tmp = np.reshape(C_new, (-1, 1))

                      #K_new = K_new + (x - c_tmp).dot((x - c_tmp).transpose())
  
                      tmp = tmp + (x - c_tmp).dot((x - c_tmp).transpose())
                  if self.constrained:
                      K_new = (tmp / X_sub.shape[0]) / np.power(np.linalg.det((tmp / X_sub.shape[0])), 1/d)
                  else:
                      K_new = tmp / X_sub.shape[0]

                  K_new = np.linalg.pinv(K_new)

                  #if end_algo and (not (self.C[k] == C_new).all()):  # If the same stop
                  #    end_algo = False

                  self.C[k] = C_new
                  self.L[k] = K_new

              diff = np.sum(np.absolute((self.C - old_c) / old_c * 100))
              if diff > tol:
                  end_algo = False
                  if debug:
                      print(f"{diff}")

              else:
                  if debug:
                      print(f"Tolerance threshold {tol} reached with diff {diff}")

                  end_algo = True

              i = i + 1

              if i > maxiter:
                  end_algo = True
                  if debug:
                      print(f"Iteration {maxiter} reached")

          return {
              "loss": loss,
              "C": self.C,
              "K": self.K,
              "L": self.L
          }