# import sys
# sys.path.append('..')

from functools import partial

import torch
import dgl

# import read_data
from gnn.Datasets import CompleteMVGraphDatasets
from gnn.networks import AutoVisualMultiGINComplete
from gnn.useful_utils import visualization_metric


def evaluate(d_loader):
    net.eval()
    with torch.no_grad():
        scs = []
        nmis = []
        losses = []
        for n_view_graphs, z, zdist, y in d_loader:
            optimizer.zero_grad()
            # inp = inp.to(torch.float32).to(device)
            # z = torch.from_numpy(np.concatenate(z)).to(torch.float32).to(device)
            n_view_graphs = n_view_graphs.to(net.device)

            z_hat = net(n_view_graphs)[0]
            # y_hat = y_hat.view(-1)
            zdist = zdist[0].to(torch.float32).to(net.device)

            # zdist = torch.stack(zdist).to(torch.float32).to(device)
            student_t_loss = net.loss_fn_kl_t(n_view_graphs, zdist)

            nmi, sc = visualization_metric.get_nmi_sc(z_hat.cpu().numpy(), y[0].tolist())

            scs.append(sc)
            nmis.append(nmi)
            losses.append(student_t_loss)

    scs_t = torch.mean(torch.Tensor(scs))
    nmis_t = torch.mean(torch.Tensor(nmis))
    losses_t = torch.mean(torch.Tensor(losses))
    net.train()
    return nmis_t, scs_t, losses_t


if __name__ == '__main__':
    device = 'cuda:0'

    # train_names = ['arrhythmia', 'wine', 'lympho', 'glass', 'vertebral',
    #                                    'wbc', 'ecoli', 'ionosphere', 'breastw', 'pima',
    #                                    'vowels',        'letter']
    # test_names = ['cardio',  'seismic', 'musk', 'speech', 'abalone']

    train_names = ['mnist_group2', 'mnist_group1', 'fmnist_group2', 'fmnist_group1', 'cifar10_group1'] #+ [f'mnist_comb{i}' for i in range(252)]
    # train_names = [f'mnist_comb{i}' for i in range(128, 252)]
    test_names = ['cifar10_group2']

    z_cali = 'none'

    precomputed_pe_path = None
    # precomputed_pe_path = '/home/****/autovisual/prepare_data/pe/pe_for_gat'

    # dummy_train_ds = GraphDatasets.DatasetGraphDataset(data_names=train_names, cdist_path='../prepare_data/clip/features',
    #                                              visual_path='../prepare_data/bo/res-2', normalize_z=True,
    #                                              z_cali_method='none',
    #                                              precomputed_pe_path='../prepare_data/pe/pe_for_gat',
    #                                              z_mu=None, z_std=None)

    train_ds = CompleteMVGraphDatasets.DatasetGraphDataset(data_names=train_names, cdist_path='/home/****/autovisual/prepare_data/clip/features',
                                                           visual_path='/home/****/autovisual/prepare_data/bo/res-2', normalize_z=True, z_cali_method=z_cali,
                                                           # z_anchor=dummy_train_ds.z[0],
                                                           precomputed_pe_path=precomputed_pe_path,
                                                           z_mu=None, z_std=None,
                                                           flip_sign_method='pos'
                                                           )




    test_ds = CompleteMVGraphDatasets.DatasetGraphDataset(data_names=test_names, cdist_path='/home/****/autovisual/prepare_data/clip/features',
                                                          visual_path='/home/****/autovisual/prepare_data/bo/res-2', normalize_z=True,
                                                          z_cali_method=z_cali,
                                                          # z_anchor=dummy_train_ds.z[0],
                                                          precomputed_pe_path=precomputed_pe_path,
                                                          z_mu=train_ds.z_mu, z_std=train_ds.z_std,
                                                          flip_sign_method='pos')

    # get_loader = partial(torch.utils.data.DataLoader, batch_size=3,
    #                                            shuffle=False,
    #                                            num_workers=4,
    #                      collate_fn=lambda x: list(zip(*x))
    #                                            )

    get_loader = partial(dgl.dataloading.GraphDataLoader, batch_size=1,
                         shuffle=False,
                         num_workers=0,
                         # collate_fn=lambda x: list(zip(*x))
                         )

    train_loader = get_loader(train_ds)
    test_loader = get_loader(test_ds)

    # a = iter(train_ds)
    # a_ = next(a)

    input_dim = train_ds[0][0].ndata['pe0'].shape[1]
    print(input_dim)

    net = AutoVisualMultiGINComplete.AutoVisualNet(input_dim=input_dim, gnn_hidden=128, gnn_out=128, out_dim=2,
                                                   n_transformer=8,
                                   n_graph_view=5, device=device)
    # f_name = './res/GNN-gat-run4-8000epoch-kl2_64dim.tar'
    # net, save = torch.load(f_name)
    net = net.to(device)

    print(net)

    n_epoch = 500

    # optimizer = torch.optim.Adam(params=net.parameters(), lr=1e-3, amsgrad=True)
    optimizer = torch.optim.AdamW(params=net.parameters(), lr=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=n_epoch)
    loss_fn = torch.nn.MSELoss()

    records = {}
    for epoch in range(n_epoch):
        running_loss = 0
        net.train()
        for n_view_graphs, z, zdist, _ in train_loader:

            optimizer.zero_grad()
            # inp = inp.to(torch.float32).to(device)
            # z = torch.from_numpy(np.concatenate(z)).to(torch.float32).to(device)
            # print(n_view_graphs.ndata['pe0'].shape)

            # zdist = zdist[0].to(torch.float32).to(device)  # only works for bs 1 so far
            # zdist = torch.stack(zdist).to(torch.float32).to(device)
            zdist = zdist.to(torch.float32).to(device)
            n_view_graphs = n_view_graphs.to(device)
            # zdist.to(torch.float32).to(device)  # only works for bs 1 so far
            # print(xs)
            # xs = (torch.ones_like(xs[0]), )
            # xs = xs[0].to(torch.float32).to(device)  # only works for bs 1 so far
            # n_view_graphs = n_view_graphs[0].to(device)  # only works for bs 1 so far
            loss = net.loss_fn_kl_t(n_view_graphs, zdist)
            # print(loss)
            # loss = net.loss_fn_kl2(xs, n_view_graphs, zdist)

            # z_hat = net(xs, n_view_graphs)
            # loss = loss_fn(z_hat, z)
            # loss = net.loss_fn_calib2(inp, y)
            # print('start opt')
            loss.backward()
            optimizer.step()
            # print('finish opt')

            # del zdist
            # del xs
            # del n_view_graphs
            # zdist = zdist.to('cpu')  # only works for bs 1 so far
            # xs = xs.to('cpu')  # only works for bs 1 so far
            # n_view_graphs = n_view_graphs.to('cpu')  # only works for bs 1 so far
            # torch.cuda.empty_cache()
            running_loss += loss.item()  # .item() is important

        scheduler.step()
        # print('+++++++++++++++++++++++++++')

        if (epoch + 1) % 10 == 0:
            # train_nmi, train_sc, train_loss = evaluate(train_loader)
            test_nmi, test_sc, test_loss = evaluate(test_loader)
            # test_mse, test_mae, test_error = evaluate(test_loader)

            # records['train_nmi'] = records.get('train_nmi', []) + [train_nmi]
            # records['train_sc'] = records.get('train_sc', []) + [train_sc]
            # records['train_loss'] = records.get('train_loss', []) + [train_loss]
            # records['val_mse'] = records.get('val_mse', []) + [val_mse.item()]
            # records['val_mae'] = records.get('val_mae', []) + [val_mae.item()]
            # records['val_error'] = records.get('val_error', []) + [val_error.item()]
            records['test_nmi'] = records.get('test_nmi', []) + [test_nmi]
            records['test_sc'] = records.get('test_sc', []) + [test_sc]
            records['test_loss'] = records.get('test_loss', []) + [test_loss]
            records['running_loss'] = records.get('running_loss', []) + [(running_loss / len(train_loader))]

            print([(k, v[-1]) for k, v in records.items()])
        if (epoch + 1) % 100 == 0:
            torch.save((net, records), f'./res/GNN-resgin2-8gt-sigma_mv-run4-{n_epoch}epoch-kl_t_64dim_svdu_complete_G_sep_gnn-epoch-{epoch}.tar')
