import os
import logging
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torch.nn.functional as F
import random
from torch.utils.data import DataLoader

from src.models.model import *
from src.other_utils.datasets import CIFAR10_truncated, FashionMNIST_truncated, MNIST_truncated
from torchvision.datasets import MNIST

logging.basicConfig()
logger = logging.getLogger()
logger.setLevel(logging.INFO)

def mkdirs(dirpath):
    try:
        os.makedirs(dirpath)
    except Exception as _:
        pass
    
    
def load_mnist_data(datadir):

    transform = transforms.Compose([transforms.ToTensor()])

    mnist_train_ds = MNIST_truncated(datadir, train=True, download=True, transform=transform)
    mnist_test_ds = MNIST_truncated(datadir, train=False, download=True, transform=transform)

    X_train, y_train = mnist_train_ds.data, mnist_train_ds.target
    X_test, y_test = mnist_test_ds.data, mnist_test_ds.target
    
    # X_train, y_train = mnist_train_ds.data, mnist_train_ds.targets
    # X_test, y_test = mnist_test_ds.data, mnist_test_ds.targets

    X_train = X_train.data.numpy()
    y_train = y_train.data.numpy()
    X_test = X_test.data.numpy()
    y_test = y_test.data.numpy()

    return (X_train, y_train, X_test, y_test)


def load_fmnist_data(datadir):

    transform = transforms.Compose([transforms.ToTensor()])

    train_ds = FashionMNIST_truncated(datadir, train=True, download=True, transform=transform)
    test_ds = FashionMNIST_truncated(datadir, train=False, download=True, transform=transform)

    X_train, y_train = train_ds.data, train_ds.target
    X_test, y_test = test_ds.data, test_ds.target

    X_train = X_train.data.numpy()
    y_train = y_train.data.numpy()
    X_test = X_test.data.numpy()
    y_test = y_test.data.numpy()

    return (X_train, y_train, X_test, y_test)


def load_cifar10_data(datadir):
    
    transform = transforms.Compose([transforms.ToTensor()])

    cifar10_train_ds = CIFAR10_truncated(datadir, train=True, download=True, transform=transform)
    cifar10_test_ds = CIFAR10_truncated(datadir, train=False, download=True, transform=transform)

    X_train, y_train = cifar10_train_ds.data, cifar10_train_ds.target
    X_test, y_test = cifar10_test_ds.data, cifar10_test_ds.target
    
    return (X_train, y_train, X_test, y_test)

def record_net_data_stats(y_train, net_dataidx_map, logdir):
    net_cls_counts = {}
    for net_i, dataidx in net_dataidx_map.items():
        unq, unq_cnt = np.unique(y_train[dataidx], return_counts=True)
        tmp = {unq[i]: unq_cnt[i] for i in range(len(unq))}
        net_cls_counts[net_i] = tmp

    logger.info('Data statistics: %s' % str(net_cls_counts))
    return net_cls_counts

def partition_data(dataset, datadir, logdir, partition, n_workers, n_sample, alpha=0.1):
    if dataset == 'mnist':
        X_train, y_train, X_test, y_test = load_mnist_data(datadir)
    elif dataset == 'fmnist':
        X_train, y_train, X_test, y_test = load_fmnist_data(datadir)
    elif dataset == 'cifar10':
        X_train, y_train, X_test, y_test = load_cifar10_data(datadir)
    else:
        raise NotImplementedError(f"Dataset {dataset} not implemented")

    n_train = y_train.shape[0]
    
    if n_workers == 1:
        net_dataidx_map = {0: np.arange(n_train)}
        return net_dataidx_map
    
    if partition == "homo":
        idxs = np.random.permutation(n_train)
        batch_idxs = np.array_split(idxs, n_workers)
        net_dataidx_map = {i: batch_idxs[i] for i in range(n_workers)}

    elif partition == "noniid-labeldir":
        min_size = 0
        min_require_size = 10
        K = 10  # classes
        if dataset in ('celeba', 'covtype', 'a9a', 'rcv1', 'SUSY'):
            K = 2

        N = y_train.shape[0]
        net_dataidx_map = {}

        while min_size < min_require_size:
            idx_batch = [[] for _ in range(n_workers)]
            for k in range(K):
                idx_k = np.where(y_train == k)[0]
                np.random.shuffle(idx_k)
                proportions = np.random.dirichlet(np.repeat(alpha, n_workers))
                ## Balance
                proportions = np.array([p * (len(idx_j) < N / n_workers) for p, idx_j in zip(proportions, idx_batch)])
                proportions = proportions / proportions.sum()
                proportions = (np.cumsum(proportions) * len(idx_k)).astype(int)[:-1]
                idx_batch = [idx_j + idx.tolist() for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))]
                min_size = min([len(idx_j) for idx_j in idx_batch])
                # if K == 2 and n_workers <= 10:
                #     if np.min(proportions) < 200:
                #         min_size = 0
                #         break

        for j in range(n_workers):
            np.random.shuffle(idx_batch[j])
            net_dataidx_map[j] = idx_batch[j]
            # if n_sample > 0 and n_sample < len(net_dataidx_map[j]):
            #     np.random.shuffle(net_dataidx_map[j])
            #     net_dataidx_map[j] = net_dataidx_map[j][:n_sample]

    elif partition > "noniid-#label0" and partition <= "noniid-#label9":
        num = eval(partition[13:])
        n_sample_c = int(n_train/(n_workers*num))  # maximal numbers of samples per class
        K = 10
        if num == 10:
            net_dataidx_map = {i:np.ndarray(0,dtype=np.int64) for i in range(n_workers)}
            for i in range(10):
                idx_k = np.where(y_train==i)[0]
                np.random.shuffle(idx_k)
                split = np.array_split(idx_k,n_workers)
                for j in range(n_workers):
                    np.random.shuffle(split[j])
                    net_dataidx_map[j]=np.append(net_dataidx_map[j],split[j][:n_sample_c])
        else:
            times=[0 for i in range(10)]
            contain=[]
            for i in range(n_workers):
                current=[i%K]
                times[i%K]+=1
                j=1
                while (j<num):
                    ind=random.randint(0,K-1)
                    if (ind not in current):
                        j=j+1
                        current.append(ind)
                        times[ind]+=1
                contain.append(current)
            net_dataidx_map ={i:np.ndarray(0,dtype=np.int64) for i in range(n_workers)}
            for i in range(K):
                idx_k = np.where(y_train==i)[0]
                np.random.shuffle(idx_k)
                split = np.array_split(idx_k,times[i])
                ids=0
                for j in range(n_workers):
                    if i in contain[j]:
                        np.random.shuffle(split[ids])  # TODO: check
                        net_dataidx_map[j]=np.append(net_dataidx_map[j],split[ids][:n_sample_c])
                        ids+=1

    else:
        raise NotImplementedError(f"Partition {partition} not implemented")

    traindata_cls_counts = record_net_data_stats(y_train, net_dataidx_map, logdir)
    # return (X_train, y_train, X_test, y_test, net_dataidx_map, traindata_cls_counts)
    return net_dataidx_map


def compute_accuracy(net, dataloader, device, dataset=None):
    net.to(device)
    
    if dataset == "synthetic":
        criterion = nn.MSELoss().to(device)
    else:
        criterion = nn.CrossEntropyLoss().to(device)
        
    was_training = False
    if net.training:
        net.eval()
        was_training = True

    correct, total, loss_sum = 0, 0, 0
    with torch.no_grad():
        for batch_idx, (x, target) in enumerate(dataloader):
            x, target = x.to(device), target.to(device)
            out = net(x)
            
            # print("out", out.shape, out) # (128, 10)
            loss = criterion(out, target)
            _, pred_label = torch.max(out.data, 1)
            
            # print("pred_label", pred_label)
            # print("target", target.data)

            total += x.data.size()[0]
            correct += (pred_label == target.data).sum().item()
            loss_sum += loss.item() * x.data.size()[0]

    if was_training:
        net.train()

    return correct/float(total), loss_sum/float(total)


def get_dataloader(dataset, datadir, train_bs, test_bs=32, dataidxs=None, net_dataidx_map=None):
    if dataset == 'mnist':
        dl_obj = MNIST_truncated
        MNIST_MEAN = 0.1307
        MNIST_STD = 0.3081
        
        transform_train = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((MNIST_MEAN,), (MNIST_STD,))
        ])
        transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((MNIST_MEAN,), (MNIST_STD,))
        ])
    
    elif dataset == 'fmnist': # TODO: no normalization
        dl_obj = FashionMNIST_truncated
        transform_train = transforms.Compose([
            transforms.ToTensor()])
        transform_test = transforms.Compose([
            transforms.ToTensor()])

    elif dataset == 'cifar10':
        dl_obj = CIFAR10_truncated

        CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
        CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]

        transform_train = transforms.Compose([
            transforms.ToPILImage(),
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
        ])

        transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
        ])

    else:
        raise NotImplementedError(f"Dataset {dataset} not implemented")
    
    # global
    train_ds = dl_obj(datadir, dataidxs=None, train=True, transform=transform_train, download=True)
    test_ds = dl_obj(datadir, train=False, transform=transform_test, download=True)
    train_dl = DataLoader(dataset=train_ds, batch_size=train_bs, shuffle=True, drop_last=False)
    test_dl = DataLoader(dataset=test_ds, batch_size=test_bs, shuffle=False, drop_last=False)
    
    train_dl_dict = {}
    if net_dataidx_map is not None:
        for idx in range(len(net_dataidx_map)):
            train_ds = dl_obj(datadir, dataidxs=net_dataidx_map[idx], train=True, transform=transform_train, download=True)
            train_dl = DataLoader(dataset=train_ds, batch_size=train_bs, shuffle=True, drop_last=False)
            train_dl_dict[idx] = train_dl
    
    return train_dl, test_dl, train_dl_dict

