import os
import numpy as np
import torch
import shutil
from torch.autograd import Variable
import json
from sklearn.metrics import confusion_matrix


class AvgrageMeter(object):

  def __init__(self):
    self.reset()

  def reset(self):
    self.avg = 0
    self.sum = 0
    self.cnt = 0

  def update(self, val, n=1):
    self.sum += val * n
    self.cnt += n
    self.avg = self.sum / self.cnt


def accuracy(output, target, topk=(1,)):
  maxk = max(topk)
  batch_size = target.size(0)

  _, pred = output.topk(maxk, 1, True, True)
  pred = pred.t()
  correct = pred.eq(target.view(1, -1).expand_as(pred))
  print(correct.shape)
  res = []
  for k in topk:
    correct_k = correct[:k].reshape(-1).float().sum(0)
    res.append(correct_k.mul_(100.0/batch_size))
  return res


class Cutout(object):
    def __init__(self, length):
        self.length = length

    def __call__(self, img):
        h, w = img.size(1), img.size(2)
        mask = np.ones((h, w), np.float32)
        y = np.random.randint(h)
        x = np.random.randint(w)

        y1 = np.clip(y - self.length // 2, 0, h)
        y2 = np.clip(y + self.length // 2, 0, h)
        x1 = np.clip(x - self.length // 2, 0, w)
        x2 = np.clip(x + self.length // 2, 0, w)

        mask[y1: y2, x1: x2] = 0.
        mask = torch.from_numpy(mask)
        mask = mask.expand_as(img)
        img *= mask
        return img


def _data_transforms_cifar10(args):
  CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
  CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]

  train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
  ])
  if args.cutout:
    train_transform.transforms.append(Cutout(args.cutout_length))

  valid_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
    ])
  return train_transform, valid_transform


def count_parameters_in_MB(model):
  return np.sum(np.prod(v.size()) for name, v in model.named_parameters() if "auxiliary" not in name)/1e6


def save_checkpoint(state, is_best, save):
  filename = os.path.join(save, 'checkpoint.pth.tar')
  torch.save(state, filename)
  if is_best:
    best_filename = os.path.join(save, 'model_best.pth.tar')
    shutil.copyfile(filename, best_filename)


def save(model, model_path):
  torch.save(model.state_dict(), model_path)


def load(model, model_path):
  model.load_state_dict(torch.load(model_path))

"""
def drop_path(x, drop_prob):
  if drop_prob > 0.:
    keep_prob = 1.-drop_prob
    mask = Variable(torch.cuda.FloatTensor(x.size(0),1, 1).bernoulli_(keep_prob))
    x.div_(keep_prob)
    x.mul_(mask)
  return x
"""
def drop_path2d(x, drop_prob):
  if drop_prob > 0.:
    keep_prob = 1.-drop_prob
    mask = Variable(torch.cuda.FloatTensor(x.size(0), 1,1, 1).bernoulli_(keep_prob))
    x.div_(keep_prob)
    x.mul_(mask)
  return x

def drop_path(x, drop_prob):
  if drop_prob > 0.:
    keep_prob = 1.-drop_prob
    mask = Variable(torch.cuda.FloatTensor(x.size(0), 1, 1).bernoulli_(keep_prob))
    x.div_(keep_prob)
    x.mul_(mask)
  return x

def create_exp_dir(path, scripts_to_save=None):
  if not os.path.exists(path):
    os.mkdir(path)
  print('Experiment dir : {}'.format(path))

  if scripts_to_save is not None:
    os.mkdir(os.path.join(path, 'scripts'))
    for script in scripts_to_save:
      dst_file = os.path.join(path, 'scripts', os.path.basename(script))
      shutil.copyfile(script, dst_file)



import pandas as pd


class BernoulliLogger:
  def __init__(self,PATH,ID): 
    if not os.path.exists("{}/{}".format(PATH,"metrics")):
      os.mkdir("{}/metrics".format(self.PATH))
    self.path = "{}/{}/{}-bin.npy".format(PATH,"metrics",ID)
  def update(self,data):
    if os.path.exists(self.path):
      x = np.load(self.path)
      np.save(self.path,np.vstack((x,data)))
    else:
      np.save(self.path,data)

class MetricLogger:
  def __init__(self,PATH):
    self.performance_data = {"ID": [],"accuracy": [], "recall":[]}
    self.PATH = PATH
    if not os.path.exists("{}/{}".format(PATH,"metrics")):
      os.mkdir("{}/metrics".format(self.PATH))
    self.initial = True 
  def load(self,ID):
    path = "{}/{}/{}".format(self.PATH,"metrics",ID)
    data = pd.read_csv(path).to_dict('list')
    self.performance_data = {"ID": data["ID"],"accuracy": data["accuracy"], "recall":data["recall"]}

  def update(self, new_data : dict):
    if self.initial and os.path.exists("{}/{}/{}".format(self.PATH,"metrics",new_data["ID"])):
      self.load(new_data["ID"])
      self.initial = False
    for i in new_data:
      self.performance_data[i].append(new_data[i])

    df = pd.DataFrame.from_dict(self.performance_data)
    df.to_csv("{}/{}/{}".format(self.PATH,"metrics",new_data["ID"]))




def calculate_train_vals(pred_tensor,gt_tensor):
    # Convert the tensors to NumPy arrays
  predictions = torch.argmax(pred_tensor, dim=1).numpy()
  labels = gt_tensor.numpy()
  # Calculate accuracy
  correct = np.sum(predictions == labels)
  total = len(labels)
  accuracy = correct / total

  # Determine the number of classes
  num_classes = int(max(labels.max(), predictions.max()) + 1)

  # Generate confusion matrix
  cm = confusion_matrix(labels, predictions)

  # Serialize the confusion matrix as a JSON string
  confusion_matrix_str = json.dumps(cm.tolist())
  return confusion_matrix_str, accuracy



