# %%
from singa import singa_wrap as singa
from singa import device
from singa import tensor
from singa import opt
from singa import layer
from singa.layer import Layer
from singa import model

import numpy as np
import time
import argparse
from PIL import Image
import os
from glob import glob
from sklearn.metrics import accuracy_score
import json
from tqdm import tqdm

from transforms import Compose, ToTensor, Normalize

np_dtype = {"float16": np.float16, "float32": np.float32}

singa_dtype = {"float16": tensor.float16, "float32": tensor.float32}

# %%
transforms = Compose([
    ToTensor(),
    Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# %%
class ClassDataset:
    def __init__(self, img_folder, transforms):
        super(ClassDataset, self).__init__()

        self.img_list = list()
        self.transforms = transforms

        classes = os.listdir(img_folder)
        for i in classes:
            images = glob(os.path.join(img_folder, i, "*"))
            for img in images:
                self.img_list.append((img, i))
    
    def __len__(self) -> int:
        return len(self.img_list)

    def __getitem__(self, index: int):
        img_path, label_str = self.img_list[index]
        img = Image.open(img_path)
        img = self.transforms.forward(img)
        label = np.array(label_str, dtype=np.int32)

        return img, label
    
    def batchgenerator(self, indexes, batch_size, data_size):
        batch_x = np.zeros((batch_size,) + data_size)
        batch_y = np.zeros((batch_size,) + (1,), dtype=np.int32)
        for idx, i in enumerate(indexes):
            sample_x, sample_y = self.__getitem__(i)
            batch_x[idx, :, :, :] = sample_x
            batch_y[idx, :] = sample_y

        return batch_x, batch_y


# %%
class Sequential:
    def __init__(self, *args):
        # super(Sequential, self).__init__()
        self.mod_list = []
        for i in args:
            self.mod_list.append(i)
    def __call__(self, input):
        for module in self.mod_list:
            input = module(input)
        return input

class CNNModel(model.Model):
    def __init__(self, num_classes):
        super(CNNModel, self).__init__()
        self.input_size = 28
        self.dimension = 4
        self.num_classes = num_classes
        
        # self.layer1 = Sequential(
        #     layer.Conv2d(16, kernel_size=3),
        #     layer.BatchNorm2d(),
        #     layer.ReLU())
        self.layer1 = layer.Conv2d(16, kernel_size=3, activation="RELU")
        self.bn1 = layer.BatchNorm2d()
        # self.layer2 = Sequential(
        #     layer.Conv2d(16, kernel_size=3),
        #     layer.BatchNorm2d(),
        #     layer.ReLU(),
        #     layer.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = layer.Conv2d(16, kernel_size=3, activation="RELU")
        self.bn2 = layer.BatchNorm2d()        
        self.pooling2 = layer.MaxPool2d(kernel_size=2, stride=2)

        # self.layer3 = Sequential(
        #     layer.Conv2d(64, kernel_size=3),
        #     layer.BatchNorm2d(),
        #     layer.ReLU())
        self.layer3 = layer.Conv2d(64, kernel_size=3, activation="RELU")
        self.bn3 = layer.BatchNorm2d()
        
        # self.layer4 = Sequential(
        #     layer.Conv2d(64, kernel_size=3),
        #     layer.BatchNorm2d(),
        #     layer.ReLU())
        self.layer4 = layer.Conv2d(64, kernel_size=3, activation="RELU")
        self.bn4 = layer.BatchNorm2d()

        # self.layer5 = Sequential(
        #     layer.Conv2d(64, kernel_size=3, padding=1),
        #     layer.BatchNorm2d(),
        #     layer.ReLU(),
        #     layer.MaxPool2d(kernel_size=2, stride=2))
        self.layer5 = layer.Conv2d(64, kernel_size=3, padding=1, activation="RELU")
        self.bn5 = layer.BatchNorm2d()
        self.pooling5 = layer.MaxPool2d(kernel_size=2, stride=2)

        self.flatten = layer.Flatten()

        # self.fc = Sequential(
        #     layer.Linear(128),
        #     layer.ReLU(),
        #     layer.Linear(128),
        #     layer.ReLU(),
        #     layer.Linear(num_classes))
        self.linear1 = layer.Linear(128)
        self.linear2 = layer.Linear(128)
        self.linear3 = layer.Linear(self.num_classes)

        self.relu = layer.ReLU()

        self.softmax_cross_entropy = layer.SoftMaxCrossEntropy()
        self.dropout = layer.Dropout(ratio=0.3)

    def forward(self, x):
        x = self.layer1(x)
        x = self.bn1(x)
        x = self.layer2(x)
        x = self.bn2(x) 
        x = self.pooling2(x)
        
        x = self.layer3(x)
        x = self.bn3(x) 
        x = self.layer4(x)
        x = self.bn4(x) 
        x = self.layer5(x)
        x = self.bn5(x) 
        x = self.pooling5(x)
        x = self.flatten(x)
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        x = self.relu(x)
        x = self.linear3(x)
        # x = self.fc(x)
        return x

    def set_optimizer(self, optimizer):
        self.optimizer = optimizer

    def train_one_batch(self, x, y, dist_option, spars):
        out = self.forward(x)
        loss = self.softmax_cross_entropy(out, y)

        if dist_option == 'plain':
            self.optimizer(loss)
        elif dist_option == 'half':
            self.optimizer.backward_and_update_half(loss)
        elif dist_option == 'partialUpdate':
            self.optimizer.backward_and_partial_update(loss)
        elif dist_option == 'sparseTopK':
            self.optimizer.backward_and_sparse_update(loss,
                                                      topK=True,
                                                      spars=spars)
        elif dist_option == 'sparseThreshold':
            self.optimizer.backward_and_sparse_update(loss,
                                                      topK=False,
                                                      spars=spars)
        return out, loss

# %%
def getACC(y_true, y_score):
    '''Accuracy metric.
    :param y_true: the ground truth labels, shape: (n_samples, n_labels) or (n_samples,) if n_labels==1
    :param y_score: the predicted score of each class,
    shape: (n_samples, n_labels) or (n_samples, n_classes) or (n_samples,) if n_labels==1 or n_classes==1
    :param task: the task of current dataset
    :param threshold: the threshold for multilabel and binary-class tasks
    '''
    y_true = y_true.squeeze()
    y_score = y_score.squeeze()

    
    ret = accuracy_score(y_true, np.argmax(y_score, axis=-1))

    return ret
def accuracy(pred, target):
    # y is network output to be compared with ground truth (int)
    y = np.argmax(pred, axis=1)
    a = (y[:,None]==target).sum()
    correct = np.array(a, "int").sum()
    return correct

# %% dataset configuration
dataset_path = "./bloodmnist"
train_path = os.path.join(dataset_path, "train")
val_path = os.path.join(dataset_path, "val") 
cfg_path = os.path.join(dataset_path, "param.json")

with open(cfg_path,'r') as load_f:
    num_class = json.load(load_f)["num_classes"]

train_dataset = ClassDataset(train_path, transforms)
val_dataset = ClassDataset(val_path, transforms)

batch_size = 256

# %% model configuration
model = CNNModel(num_classes=num_class)
criterion = layer.SoftMaxCrossEntropy()
# optimizer_ft = opt.SGD(lr=0.005, momentum=0.9, weight_decay=1e-5, dtype=singa_dtype["float32"])
optimizer_ft = opt.Adam(lr=1e-3)
# optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr)
# lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_ft, mode='max', patience=5, threshold=1e-3)

# %% start training
dev = device.create_cpu_device()
dev.SetRandSeed(0)
np.random.seed(0)

tx = tensor.Tensor(
        (batch_size, 3, model.input_size, model.input_size), dev,
        singa_dtype['float32'])
ty = tensor.Tensor((batch_size,), dev, tensor.int32)

num_train_batch = train_dataset.__len__() // batch_size
num_val_batch = val_dataset.__len__() // batch_size
idx = np.arange(train_dataset.__len__(), dtype=np.int32)

model.set_optimizer(optimizer_ft)
model.compile([tx], is_train=True, use_graph=False, sequential=False)
dev.SetVerbosity(0)

max_epoch = 100
for epoch in range(max_epoch):
    print('start**************************************************')
    
    start_time = time.time()
    # np.random.shuffle(idx)

    train_correct = np.zeros(shape=[1], dtype=np.float32)
    test_correct = np.zeros(shape=[1], dtype=np.float32)
    train_loss = np.zeros(shape=[1], dtype=np.float32)

    model.train()
    for b in tqdm(range(num_train_batch)):
        # forward + backward + optimize

        x, y = train_dataset.batchgenerator(idx[b * batch_size:(b + 1) * batch_size], 
            batch_size=batch_size, data_size=(3, model.input_size, model.input_size))
        x = x.astype(np_dtype['float32'])
        # print(tx.size())
        # print(y.dtype)

        tx.copy_from_numpy(x)
        ty.copy_from_numpy(y)

        out, loss = model(tx, ty, dist_option="plain", spars=None)
        # print(out)
        # print(loss)
        # train_correct += getACC(y, tensor.to_numpy(out))
        train_correct += accuracy(tensor.to_numpy(out), y)
        train_loss += tensor.to_numpy(loss)[0]
    print('Training loss = %f, training accuracy = %f' %
                  (train_loss, train_correct /
                   (num_train_batch * batch_size)))
    model.eval()

    for b in tqdm(range(num_val_batch)):
        # forward + backward + optimize
        x, y = train_dataset.batchgenerator(idx[b * batch_size:(b + 1) * batch_size], 
            batch_size=batch_size, data_size=(3, model.input_size, model.input_size))
        x = x.astype(np_dtype['float32'])

        tx.copy_from_numpy(x)
        ty.copy_from_numpy(y)

        out = model(tx)

        # test_correct += getACC(y, tensor.to_numpy(out))
        test_correct += accuracy(tensor.to_numpy(out), y)
    
    print('Evaluation accuracy = %f, Elapsed Time = %fs' %
                  (test_correct / (num_val_batch * batch_size),
                   time.time() - start_time))