text: """
  pynet: self supervised clustering
  =================================
  
  Credit: A Grigis
  """
  
  import os
  import sys
  if "CI_MODE" in os.environ:
      sys.exit()
  
  # Imports
  import collections
  import logging
  import pynet
  from pynet.metrics import SKMetrics
  from pynet.datasets import DataManager
  from pynet.interfaces import DeepLearningInterface
  from pynet.interfaces import DeepClusterClassifier
  from pynet.models import BrainNetCNN
  from pynet.utils import setup_logging
  from pynet.plotting import Board, update_board
  from pynet.models.deepcluster import update_pseudo_labels
  import torch
  import torch.nn as nn
  from torch.utils.data.sampler import Sampler
  import numpy as np
  import matplotlib.pyplot as plt
  import seaborn as sns
  import pandas as pd
  from sklearn.cluster import MiniBatchKMeans, KMeans
  from sklearn.metrics import classification_report
  from sklearn.metrics import roc_curve, auc
  try:
      import faiss
  except:
      pass
  
  
  # Global Parameters
  OUTDIR = "/tmp/graph_connectome"
  BATCH_SIZE = 5
  N_EPOCHS = 20
  N_CLUSTERS = 2
  N_SAMPLES = 40
  AVOID_EMPTY_CLUSTERS = False
  UNIFORM_SAMPLING = True
  setup_logging(level="info")
  
  
  # Load the data
  data = []
  labels = []
  for idx in range(N_CLUSTERS):
  
      x_train = np.ones((N_SAMPLES, 1, 90, 90)) * idx
      x_train += (np.random.rand(*x_train.shape) - 0.5) * 0.01
      y_train = np.asarray([idx] * N_SAMPLES)
      print("sub data: x {0} - y {1}".format(x_train.shape, y_train.shape))
      data.append(x_train)
      labels.extend([idx] * len(x_train))
  
      # Show example noisy training data that have the signatures applied.
      # It's not obvious to the human eye the subtle differences, but the cross
      # row and column above perturbed the below matrices with the y weights.
      # Show in the title how much each signature is weighted by.
      plt.figure(figsize=(16, 4))
      for idx in range(3):
          plt.subplot(1, 3, idx + 1)
          plt.imshow(np.squeeze(x_train[idx]), interpolation="None")
          plt.colorbar()
          plt.title(y_train[idx])
  
  data = np.concatenate(data, axis=0)
  labels = np.asarray(labels)
  print("dataset: x {0} - y {1}".format(data.shape, labels.shape))
  
  
  class UniformLabelSampler(Sampler):
      """ Samples elements uniformely accross pseudo labels.
      """
      def __init__(self, data_array):
          """ Init class.
  
          Parameters
          ----------
          data_array: ArrayDataset
              the train data array that contains the pseudo labels.
          """
          self.n_samples = len(data_array)
          self.data_array = data_array
          self.indexes = self.generate_indexes_epoch()
  
      def generate_indexes_epoch(self):
          """ Generate sampling indexes.
          """
          labels = self.data_array.labels
          clusters_to_images = self.get_clusters(labels)
          n_non_empty_clusters = len(clusters_to_images)
          size_per_pseudolabel = int(self.n_samples / n_non_empty_clusters) + 1
          res = np.array([])
          for name, cluster_indexes in clusters_to_images.items():
              indexes = np.random.choice(
                  cluster_indexes, size_per_pseudolabel,
                  replace=(len(cluster_indexes) <= size_per_pseudolabel))
              res = np.concatenate((res, indexes))
          np.random.shuffle(res)
          res = list(res.astype("int"))
          if len(res) >= self.n_samples:
              return res[:self.n_samples]
          res += res[: (self.n_samples - len(res))]
          return res
  
      def get_clusters(self, labels):
          """ Get indexes associated to each cluster.
          """
          tally = collections.defaultdict(list)
          for idx, item in enumerate(labels):
              tally[item].append(idx)
          return tally
  
      def __iter__(self):
          return iter(self.indexes)
  
      def __len__(self):
          return len(self.indexes)
  
  
  # Create data manager
  if UNIFORM_SAMPLING:
      sampler = UniformLabelSampler
  else:
      sampler = "random"
  manager = DataManager.from_numpy(
      train_inputs=data, train_labels=np.zeros(labels.shape),
      batch_size=BATCH_SIZE, sampler=sampler)
  
  
  class FKmeans(object):
      def __init__(self, n_clusters):
          self.n_clusters = n_clusters
  
      def fit(self, data):
          n_data, d = data.shape
          self.clus = faiss.Kmeans(d, self.n_clusters)
          self.clus.seed = np.random.randint(1234)
          self.clus.niter = 20
          self.clus.max_points_per_centroid = 10000000
          self.clus.train(data)
  
      def predict(self, data):
          _, I = self.clus.index.search(data, 1)
          losses = self.clus.obj
          print("k-means loss evolution: {0}".format(losses))
          return np.asarray([int(n[0]) for n in I])
  
  
  def my_loss(x, y):
      criterion = nn.CrossEntropyLoss()
      print("  x: {0} - {1}".format(x.shape, x.dtype))
      print(torch.argmax(x, dim=1))
      print("  y: {0} - {1}".format(y.shape, y.dtype))
      print(y)
      return criterion(x, y)
  
  
  # Create model
  train_loader = manager.get_dataloader(train=True, fold_index=0).train
  if AVOID_EMPTY_CLUSTERS:
      kmeans = FKmeans(n_clusters=N_CLUSTERS)
  else:
      kmeans = KMeans(
          n_clusters=N_CLUSTERS,
          random_state=None,
          # verbose=100,
          max_iter=20)
  net = BrainNetCNN(
      input_shape=(90, 90),
      in_channels=1,
      num_classes=N_CLUSTERS,
      nb_e2e=32,
      nb_e2n=64,
      nb_n2g=30,
      dropout=0,
      leaky_alpha=0.1,
      twice_e2e=False,
      dense_sml=False)
  net_params = pynet.NetParameters(
      network=net,
      clustering=kmeans,
      data_loader=train_loader,
      n_batchs=10,
      pca_dim=6,
      assignment_logfile=None,
      use_cuda=False)
  model = DeepClusterClassifier(
      net_params,
      optimizer_name="SGD",
      learning_rate=0.001,
      momentum=0.9,
      weight_decay=10**-5,
      # loss=my_loss)
      loss_name="CrossEntropyLoss")
  model.board = Board(port=8097, host="http://localhost", env="deepcluster")
  model.add_observer("before_epoch", update_pseudo_labels)
  model.add_observer("after_epoch", update_board)
  
  
  # Train model
  test_history, train_history = model.training(
      manager=manager,
      nb_epochs=N_EPOCHS,
      checkpointdir=None,
      fold_index=0,
      scheduler=None,
      with_validation=False)
  
  
  # Test model
  manager = DataManager.from_numpy(
      test_inputs=data, test_labels=labels, batch_size=BATCH_SIZE)
  test_model = DeepLearningInterface(
      model=model.model.network,
      optimizer_name="SGD",
      learning_rate=0.01,
      momentum=0.9,
      weight_decay=10**-5,
      loss_name="CrossEntropyLoss")
  y_pred, X, y_true, loss, values = test_model.testing(
      manager=manager,
      with_logit=True,
      # logit_function="sigmoid",
      predict=False)
  print(y_pred.shape, X.shape, y_true.shape)
  
  
  # Inspect results
  result = pd.DataFrame.from_dict(collections.OrderedDict([
      ("pred", (np.argmax(y_pred, axis=1)).astype(int)),
      ("truth", y_true.squeeze()),
      ("prob_0", y_pred[:, 0].squeeze()),
      ("prob_1", y_pred[:, 1].squeeze())]))
  print(result)
  y_pred = np.argmax(y_pred, axis=1)
  fig, ax = plt.subplots()
  cmap = plt.get_cmap("Blues")
  cm = SKMetrics("confusion_matrix", with_logit=False)(y_pred, y_true)
  sns.heatmap(cm, cmap=cmap, annot=True, fmt="g", ax=ax)
  ax.set_xlabel("predicted values")
  ax.set_ylabel("actual values")
  print(classification_report(y_true, y_pred >= 0.4))
  fpr, tpr, _ = roc_curve(y_true, y_pred)
  roc_auc = auc(fpr, tpr)
  plt.figure()
  plt.plot(fpr, tpr, color="darkorange", lw=2,
           label="ROC curve (area = %0.2f)" % roc_auc)
  plt.plot([0, 1], [0, 1], color="navy", lw=2, linestyle="--")
  plt.xlim([0.0, 1.0])
  plt.ylim([0.0, 1.05])
  plt.xlabel("False Positive Rate")
  plt.ylabel("True Positive Rate")
  plt.title("Receiver operating characteristic")
  plt.legend(loc="lower right")
  
  plt.show()
  