from torch_geometric import __version__ as pyg_version

from torch_geometric.loader import DataLoader

from torch import __version__ as torch_version

from torch import FloatTensor, max, save, load

from torch.cuda import is_available, set_device, empty_cache, device_count

from torch.distributed import init_process_group, barrier

from torch.distributed.optim.zero_redundancy_optimizer import ZeroRedundancyOptimizer

from torch.utils.data.distributed import DistributedSampler

from torch.nn.parallel import DistributedDataParallel as DDP

from torch.distributed.elastic.multiprocessing.errors import record

from torch.utils.tensorboard.writer import SummaryWriter

from torch.optim import Adam

from torch.nn import CrossEntropyLoss

from pandas import read_csv

from os import environ

from tqdm import tqdm

from numpy import array, isin, argwhere

from Network import CHD_GNN

from Utilities import CHD_Dataset, __Load_Adjacency__

from Metrics import Accuracy_Util

@record

def main():

    DIRECTORY = '/home/sojo/Documents/ImageCHD/ImageCHD_dataset/'

    EPOCHS = 200

    init_process_group('nccl')

    environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

    local_rank = int(environ['LOCAL_RANK'])

    global_rank = int(environ['RANK'])

    batch_size = int(environ['WORLD_SIZE'])

    if global_rank == 0:

        print('PyG version: ', pyg_version)

        print('Torch version: ', torch_version)

        print('GPU available: ', is_available())

        print('GPU count: ', device_count())

    set_device(local_rank)

    empty_cache()

    model = CHD_GNN().to('cuda:' + str(local_rank))

    model = DDP(model, device_ids = [local_rank])

    adjacency = __Load_Adjacency__(DIRECTORY + 'ADJACENCY/')

    train_metadata = read_csv(filepath_or_buffer = DIRECTORY + 'train_dataset_info.csv')

    train_dataset = CHD_Dataset(metadata = train_metadata, directory = DIRECTORY, adjacency = adjacency)

    train_sampler = DistributedSampler(train_dataset, shuffle = True)

    train_dataloader = DataLoader(dataset = train_dataset, batch_size = 4,

                                num_workers = 8, persistent_workers = True,

                                sampler = train_sampler, pin_memory = True,

                                prefetch_factor = 4)

    eval_metadata = read_csv(filepath_or_buffer = DIRECTORY + 'eval_dataset_info.csv')

    eval_dataset = CHD_Dataset(metadata = eval_metadata, directory = DIRECTORY, adjacency = adjacency)

    eval_sampler = DistributedSampler(eval_dataset, shuffle = True)

    eval_dataloader = DataLoader(dataset = eval_dataset, batch_size = 4,

                                num_workers = 8, persistent_workers = True,

                                sampler = eval_sampler, pin_memory = True,

                                prefetch_factor = 4)

    loss_module = CrossEntropyLoss(weight = FloatTensor([48538./38830., 48538./1387.,

                                                        48538./1387., 48538./1387.,

                                                        48538./1387., 48538./1387.,

                                                        48538./1387., 48538./1386.]).to(local_rank))

    optimizer = ZeroRedundancyOptimizer(model.parameters(), Adam, amsgrad = True)

    writer = SummaryWriter('GNN_Experiment')

    for epoch in range(EPOCHS):

        train_metrics = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,

                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0])

        eval_metrics = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,

                              0, 0, 0, 0, 0, 0, 0, 0, 0, 0])

        train_loss = 0.0

        eval_loss = 0.0

        if global_rank == 0:

            print('--------------------------------------------------')

            print('Epoch ', epoch + 1)

        model.train()

        train_sampler.set_epoch(epoch = epoch + 1)

        eval_sampler.set_epoch(epoch = epoch + 1)

        barrier()

        for batch in tqdm(train_dataloader, total = len(train_dataloader)):

            optimizer.zero_grad()

            batch = batch.to(local_rank)

            preds = model(batch.x, batch.edge_index)

            _, pred_labels = max(preds, dim = 1)

            loss = loss_module(preds, batch.y)

            loss.backward()

            optimizer.step()

            metrics = Accuracy_Util(batch.y, pred_labels)

            mask = argwhere(isin(metrics, -1)^True)

            train_metrics[mask] += metrics[mask]

            mask += 10

            train_metrics[mask] += 1

        model.eval()

        for batch in tqdm(eval_dataloader, total = len(eval_dataloader)):

            batch = batch.to(local_rank)

            preds = model(batch.x, batch.edge_index)

            _, pred_labels = max(preds, dim = 1)

            loss = loss_module(preds, batch.y)

            metrics = Accuracy_Util(batch.y, pred_labels)

            mask = argwhere(isin(metrics, -1)^True)

            eval_metrics[mask] += metrics[mask]

            mask += 10

            eval_metrics[mask] += 1

        if global_rank == 0:

            print('-----TRAINING METRICS-----')

            print('Loss: ', train_loss / float(len(train_dataloader)))

            writer.add_scalar('train_loss', train_loss / float(len(train_dataloader)),

                              global_step = epoch + 1)

            for i in range(0, 10):

                print('Accuracy ', i, ': ', train_metrics[i] / float(train_metrics[i+10]))

                writer.add_scalar('train_accuracy_' + str(i), train_metrics[i] / float(train_metrics[i+10]),

                                  global_step = epoch + 1)

            print('-----EVALUATION METRICS-----')

            print('Loss: ', eval_loss / float(len(eval_dataloader)))

            writer.add_scalar('eval_loss', eval_loss / float(len(eval_dataloader)),

                              global_step = epoch + 1)

            for i in range(0, 10):

                print('Accuracy ', i, ': ', eval_metrics[i] / float(eval_metrics[i+10]))

                writer.add_scalar('eval_accuracy_' + str(i), eval_metrics[i] / float(eval_metrics[i+10]),

                                  global_step = epoch + 1)

            checkpoint_path = 'MODELS/gnn_' + str(epoch + 1) + '.checkpoint'

            save(model.module.state_dict(), checkpoint_path)

        barrier()

        map_location = {'cuda:%d' % 0: 'cuda:%d' % local_rank}

        model.load_state_dict(load(checkpoint_path, map_location = map_location))

if __name__ == '__main__':

    main()