import argparse
import os
import pickle

import numpy as np
import torch
import torch.nn as nn
from ogb.lsc import MAG240MDataset
from sklearn.metrics import accuracy_score
from sklearn.model_selection import KFold
from torch.optim import Adam
from tqdm import tqdm

parser = argparse.ArgumentParser()
parser.add_argument('--commit', type=str, default='rgat')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--root', type=str, default='.')
parser.add_argument('--seed', type=int, default=2021)
parser.add_argument('--hidden_channels', type=int, default=2048)
parser.add_argument('--out_channels', type=int, default=153)
args = parser.parse_args()


class Model(nn.Module):

    def __init__(self, args):
        super().__init__()

        hidden_channels = args.hidden_channels
        out_channels = args.out_channels
        self.mlp = nn.Sequential(
            nn.Linear(hidden_channels, hidden_channels),
            nn.BatchNorm1d(hidden_channels),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(hidden_channels, out_channels),
        )

    def forward(self, inps):
        logits = self.mlp(inps)
        return logits


commit = args.commit
device = torch.device(f'cuda:{args.device}')

# set seeds
seed = args.seed
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

if not os.path.exists('./saved'):
    os.makedirs('./saved/random', exist_ok=True)
    os.makedirs('./saved/finetune', exist_ok=True)

if not os.path.exists('./post_results'):
    os.makedirs('./post_results', exist_ok=True)

dataset = MAG240MDataset(root=args.root)

split_dict = dataset.get_idx_split()
valid_idx = split_dict['valid']  # numpy array storing indices of validation paper nodes
test_idx = split_dict['test']

paper_label = dataset.paper_label
valid_label = paper_label[valid_idx]
test_label = paper_label[test_idx]

train_X = np.load(f'./results/{commit}/valid_preds.npy')
train_Y = valid_label

test_X = np.load(f'./results/{commit}/test_preds.npy')

mlp_parm = torch.load(f'./results/{commit}/mlp.pt', map_location=lambda storage, loc: storage.cuda(args.device))
# k-fold
kf = KFold(n_splits=5, shuffle=False)  # init KFold
train_X = torch.from_numpy(train_X).float().to(device)
train_Y = torch.from_numpy(train_Y).long().to(device)

test_X = torch.from_numpy(test_X).float().to(device)

# base result
model = Model(args)
model.to(device)
model.mlp.load_state_dict(mlp_parm)
model.eval()
with torch.no_grad():
    logits = model(test_X)
    logits = torch.softmax(logits, dim=1)
base_test_logits = logits

with torch.no_grad():
    logits = model(train_X)
    logits = torch.softmax(logits, dim=1)
base_valid_logits = logits

base_acc = accuracy_score(train_Y.cpu().numpy(), (base_valid_logits).argmax(1).cpu().numpy())
print(f'base valid acc: {base_acc}')

# finetune result
for idx, (train_index, valid_index) in enumerate(kf.split(train_X)):
    model = Model(args)
    model.to(device)
    model.mlp.load_state_dict(mlp_parm)
    optim = Adam(model.parameters(), lr=0.0001, weight_decay=1e-5)
    ce = nn.CrossEntropyLoss()

    train_inps = train_X[train_index]
    train_labs = train_Y[train_index]

    valid_inps = train_X[valid_index]
    valid_labs = train_Y[valid_index]

    pbar = tqdm(range(1000))
    best_acc = 0
    for i in pbar:
        model.train()
        optim.zero_grad()
        logits = model(train_inps)
        loss = ce(logits, train_labs)
        loss.backward()
        optim.step()

        if i % 1 == 0:
            model.eval()
            with torch.no_grad():
                logits = model(valid_inps).cpu().numpy()
            valid_acc = accuracy_score(valid_labs.cpu().numpy(), logits.argmax(1))
            if valid_acc > best_acc:
                torch.save(model.state_dict(), f'./saved/finetune/{commit}_{idx}.pt')
            best_acc = max(valid_acc, best_acc)
            pbar.set_postfix(valid_acc=valid_acc, best_acc=best_acc)

finetune_test_logits = None
finetune_valid_logits = torch.zeros(train_X.shape[0], args.out_channels).to(device)
for idx, (train_index, valid_index) in enumerate(kf.split(train_X)):
    mlp_parm = torch.load(f'./saved/finetune/{commit}_{idx}.pt')
    model = Model(args)
    model.to(device)
    model.load_state_dict(mlp_parm)

    model.eval()

    with torch.no_grad():
        logits = model(train_X[valid_index])
        logits = torch.softmax(logits, dim=1)
    finetune_valid_logits[valid_index] = logits

    with torch.no_grad():
        logits = model(test_X)
        logits = torch.softmax(logits, dim=1)
    finetune_test_logits = logits if finetune_test_logits is None else logits + finetune_test_logits

finetune_acc = accuracy_score(train_Y.cpu().numpy(), (finetune_valid_logits).argmax(1).cpu().numpy())
print(f'finetune valid acc: {finetune_acc}')

# random result
for idx, (train_index, valid_index) in enumerate(kf.split(train_X)):  # 调用split方法切分数据
    model = Model(args)
    model.to(device)
#     model.mlp.load_state_dict(mlp_parm)
    optim = Adam(model.parameters(), lr=0.0001, weight_decay=1e-5)
    ce = nn.CrossEntropyLoss()

    model.to(device)
    train_inps = train_X[train_index]
    train_labs = train_Y[train_index]

    valid_inps = train_X[valid_index]
    valid_labs = train_Y[valid_index]

    pbar = tqdm(range(600))
    best_acc = 0
    for i in pbar:
        model.train()
        optim.zero_grad()
        logits = model(train_inps)
        loss = ce(logits, train_labs)
        loss.backward()
        optim.step()

        if i % 1 == 0:
            model.eval()
            with torch.no_grad():
                logits = model(valid_inps).cpu().numpy()
            valid_acc = accuracy_score(valid_labs.cpu().numpy(), logits.argmax(1))
            if valid_acc > best_acc:
                torch.save(model.state_dict(), f'./saved/random/{commit}_{idx}.pt')
            best_acc = max(valid_acc, best_acc)
            pbar.set_postfix(valid_acc=valid_acc, best_acc=best_acc)

random_test_logits = None
random_valid_logits = torch.zeros(train_X.shape[0], args.out_channels).to(device)
for idx, (train_index, valid_index) in enumerate(kf.split(train_X)):
    mlp_parm = torch.load(f'./saved/random/{commit}_{idx}.pt')
    model = Model(args)
    model.to(device)
    model.load_state_dict(mlp_parm)

    model.eval()

    with torch.no_grad():
        logits = model(train_X[valid_index])
        logits = torch.softmax(logits, dim=1)
    random_valid_logits[valid_index] = logits

    with torch.no_grad():
        logits = model(test_X)
        logits = torch.softmax(logits, dim=1)
    random_test_logits = logits if random_test_logits is None else logits + random_test_logits

random_acc = accuracy_score(train_Y.cpu().numpy(), (random_valid_logits).argmax(1).cpu().numpy())
print(f'random valid acc: {random_acc}')

# save result
result_dict = {
    'valid': {
        'base': base_valid_logits,
        'random': random_valid_logits,
        'finetune': finetune_valid_logits
    },
    'test': {
        'base': base_test_logits,
        'random': random_test_logits,
        'finetune': finetune_test_logits
    }
}


with open(f'./post_results/{commit}_result.pkl', 'wb') as f:
    pickle.dump(result_dict, f)
