from argparse import ArgumentParser

from itertools import chain

import os

import sys

from sklearn.model_selection import train_test_split

import torch

import argcfg

from coldgnn.configs.coldgnn_default_config import load_coldgnn_default_config, ColdGNNConfig

import pickle

from coldgnn.layers.transformer import load_vit_weights_to_transformers

from coldgnn.utils.data_loader_utils import create_index_dataloader, create_tensor_dataloader

from tqdm import tqdm

from transformers import get_linear_schedule_with_warmup

from transformers import get_constant_schedule_with_warmup

# poly

from transformers import get_polynomial_decay_schedule_with_warmup

from transformers import get_cosine_with_hard_restarts_schedule_with_warmup

from transformers import get_cosine_schedule_with_warmup

from coldgnn.utils.graph_sampling_utils import sample_graph, sample_graph_by_edge_ratio

import json

import shutil

config_name = None

use_echoless_feat = True

target_dtype = torch.float16

# target_dtype = torch.float32

use_pre_train = False

device = "cuda"

parser = ArgumentParser()

parser.add_argument('--dataset', type=str, required=True)

parser.add_argument('--seed', type=int, required=True)

parser.add_argument("--gpu", type=int, required=True)

config_class = ColdGNNConfig

parser = argcfg.add_args_by_config_class(parser, config_class)

args = parser.parse_args()

dataset_name = args.dataset

os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

from coldgnn.layers.ntsformer import NTSFormer

from coldgnn.pre_compute import create_target_h_list

from coldgnn.utils.nested_data_loader_utils import NestedDataLoader

from coldgnn.utils.dgl_utils import process_block_with_self_loops, aggregate_neighbors

from coldgnn.utils.nested_data_utils import nested_gather, nested_map

from coldgnn.utils.torch_utils import count_parameters

from dataclasses import asdict

from coldgnn.datasets.load_data import load_dgl_data, update_graph_or_config_or_split

import sys

import torch

import torch.nn.functional as F

import torchmetrics

import shortuuid

import time

import numpy as np

import dgl

from coldgnn.utils.random_utils import reset_seed

reset_seed(args.seed)

start_time = time.time()

run_id = shortuuid.uuid()

g, features, labels, num_classes, multi_label, (train_index, valid_index, test_index), (train_mask, valid_mask, test_mask), (t_channels, v_channels) = \

load_dgl_data(dataset_name, use_pre_train=use_pre_train, device=device, config_name=config_name)

num_total_nodes = g.num_nodes()

config = load_coldgnn_default_config(dataset_name, use_pre_train=use_pre_train, config_name=config_name, use_echoless_feat=use_echoless_feat)

config = argcfg.combine_args_into_config(config, args)

print(config)

g, config, features, labels, train_index, valid_index, test_index, train_mask, valid_mask, test_mask, \

(valid_text_miss_index, valid_visual_miss_index, valid_no_miss_index), \

(test_text_miss_index, test_visual_miss_index, test_no_miss_index) = \

update_graph_or_config_or_split(args, g, config, features, labels,

train_index, valid_index, test_index, train_mask, valid_mask, test_mask,

t_channels, v_channels)

model_name = config.model_name

result_dir = config.output_dir

if not os.path.exists(result_dir):

os.makedirs(result_dir, exist_ok=True)

result_path = os.path.join(result_dir, "{}.json".format(run_id))

tmp_result_path = os.path.join(result_dir, "{}.json.tmp".format(run_id))

use_pos_enc = False

pre_method = config.pre_method

raw_g = g

def adamw_optimizer_create_func(model):

params_with_decay = []

params_no_decay = []

for name, param in model.named_parameters():

if not param.requires_grad:

continue # Skip frozen parameters

if "bias" in name or "LayerNorm.weight" in name:

params_no_decay.append(param) # No weight decay for biases & LayerNorm

else:

params_with_decay.append(param) # Apply weight decay to other params

optimizer = torch.optim.AdamW([

{"params": params_with_decay, "lr": config.lr, "weight_decay": 0.01},

{"params": params_no_decay, "lr": config.lr, "weight_decay": 0.0},

])

return optimizer

callbacks = []

def train_end_to_end(model):

model_is_pre_compute = True

for callback in callbacks:

callback.model = model

for callback in callbacks:

callback.on_train_begin(config)

eval_batch_size = config.batch_size * 2

# if num_gnn_layers > 2 and dataset_name not in ["ogbn-arxiv"]:

train_batch_size = config.batch_size

time_dict = {

"start": time.time()

}

time_dict["pre_compute"] = time_dict["start"]

autocast_device = "cuda"

autocast_dtype = torch.float32

loss_func = torch.nn.CrossEntropyLoss(reduction="none")

print("model:\n", model)

print("num parameters:", sum(p.numel() for p in model.parameters()))

optimizer = adamw_optimizer_create_func(model)

num_warmup_steps = 10

# scheduler = None

scheduler = get_linear_schedule_with_warmup(optimizer,

num_warmup_steps=num_warmup_steps,

num_training_steps=config.num_epochs)

best_valid_score = -100000.0

early_stop_epoch = 0

early_stop_train_scores = None

early_stop_valid_scores = None

early_stop_test_scores = None

patience_counter = 0

should_stop = False

non_valid_test_mask = (~valid_mask) & (~test_mask)

non_valid_test_index = torch.arange(g.num_nodes(), device=device)[non_valid_test_mask]

labelled_labels = labels[labelled_all_index]

def evaluate_end_to_end():

model.eval()

f1_macro = torchmetrics.F1Score(task="multiclass", num_classes=int(num_classes), average="macro").to(device)

with torch.no_grad():

with torch.autocast(device_type=autocast_device, dtype=autocast_dtype):

batch_logits_list = []

for batch_node_index, _ in tqdm(labelled_all_data_loader):

batch_target_h_list_list = nested_map(target_h_list_list, lambda x: x[batch_node_index.to(x.device)].to(device))

batch_logits = model(batch_target_h_list_list)

batch_logits = batch_logits.detach().cpu()

batch_logits_list.append(batch_logits)

logits = torch.cat(batch_logits_list, dim=0)

del batch_logits_list

torch.cuda.empty_cache() # Clear unused memory

logits = logits.to(device)

y_pred = logits.argmax(dim=-1)

corrects = (y_pred == labelled_labels).float()

train_acc = corrects[labelled_virtual_train_index].mean().item()

valid_acc = corrects[labelled_virtual_valid_index].mean().item()

test_acc = corrects[labelled_virtual_test_index].mean().item()

# valid_text_miss_acc = corrects[labelled_virtual_text_miss_index].mean().item()

# valid_visual_miss_acc = corrects[labelled_virtual_visual_miss_index].mean().item()

# valid_no_miss_acc = corrects[labelled_virtual_no_miss_index].mean().item()

test_text_miss_acc = corrects[labelled_virtual_text_miss_index].mean().item()

test_visual_miss_acc = corrects[labelled_virtual_visual_miss_index].mean().item()

test_no_miss_acc = corrects[labelled_virtual_no_miss_index].mean().item()

f1_macro_scores = []

for split_index in [labelled_virtual_train_index, labelled_virtual_valid_index, labelled_virtual_test_index]:

f1_macro.reset()

f1_macro.update(y_pred[split_index], labelled_labels[split_index])

f1_macro_scores.append(f1_macro.compute().item())

train_f1_macro, valid_f1_macro, test_f1_macro = f1_macro_scores

train_scores = {

"accuracy": train_acc,

"f1_macro": train_f1_macro

}

valid_scores = {

"accuracy": valid_acc,

"f1_macro": valid_f1_macro,

}

test_scores = {

"accuracy": test_acc,

"f1_macro": test_f1_macro,

"text_miss_accuracy": test_text_miss_acc,

"visual_miss_accuracy": test_visual_miss_acc,

"no_miss_accuracy": test_no_miss_acc

}

return train_scores, valid_scores, test_scores

# use_processed_blocks = False

def create_next_train_edge_batch_generator():

while True:

for batch_edge_index in train_edge_data_loader:

yield batch_edge_index

next_train_edge_batch_generator = create_next_train_edge_batch_generator()

for epoch in range(1, config.num_epochs + 1):

model.train()

pbar = tqdm(train_data_loader)

for step, batch_data in enumerate(pbar):

if config.num_max_steps is not None and step >= config.num_max_steps:

break

with torch.autocast(device_type=autocast_device, dtype=autocast_dtype):

# batch_labels = batch_labels.to(device)

def forward_func(batch_node_index,

**extra_kwargs):

def forward_by_node_index(batch_node_index):

batch_target_h_list_list = nested_map(target_h_list_list, lambda x: x[batch_node_index.to(x.device)].to(device))

return model(batch_target_h_list_list, **extra_kwargs)

output = forward_by_node_index(batch_node_index)

# batch_target_h_list_list = nested_map(target_h_list_list, lambda x: x[batch_node_index.to(x.device)].to(device))

# logits, global_h = model(batch_target_h_list_list, return_all=True)

return output

# losses = loss_func(logits, batch_labels)

# loss = losses.mean()

batch_node_index, batch_labels, batch_train_mask = batch_data

logits, loss = model.nts_forward_and_compute_loss(epoch, batch_node_index, batch_labels, batch_train_mask, forward_func, config)

optimizer.zero_grad()

loss.backward()

optimizer.step()

with torch.no_grad():

batch_acc = (logits.argmax(dim=-1) == batch_labels).float().mean().item()

pbar.set_postfix({

"epoch": epoch,

"loss": loss.item(),

"acc": batch_acc

})

for callback in callbacks:

callback.on_epoch_end(epoch, config, logs=None)

if scheduler is not None:

scheduler.step()

print("current learning_rate: ", scheduler.get_last_lr())

if epoch % config.validation_freq == 0:

train_scores, valid_scores, test_scores = evaluate_end_to_end()

time_dict["train"] = time.time()

early_stop_metric_name = "accuracy"

valid_score = valid_scores[early_stop_metric_name]

if valid_score > best_valid_score:

best_valid_score = valid_score

early_stop_train_scores = train_scores

early_stop_valid_scores = valid_scores

early_stop_test_scores = test_scores

early_stop_epoch = epoch

patience_counter = 0

else:

patience_counter += config.validation_freq

if patience_counter >= config.patience:

should_stop = True

train_time = time_dict["train"] - time_dict["pre_compute"]

all_time = time_dict["train"] - time_dict["start"]

combined_config_dict = vars(args)

for k, v in asdict(config).items():

combined_config_dict[k] = v

result_dict = {

"epoch": epoch,

# "loss": loss.detach().cpu().item(),

"patience": patience_counter,

"early_stop_epoch": early_stop_epoch,

**{f"train_{k}": v for k, v in early_stop_train_scores.items()},

**{f"val_{k}": v for k, v in early_stop_valid_scores.items()},

**early_stop_test_scores,

"pre_compute_time": 0.0,

"train_time": train_time,

"all_time": all_time

}

# to demical 4

print_result_dict = {k: round(v, 4) if isinstance(v, float) else v for k, v in result_dict.items()}

print("patience_counter = {}".format(patience_counter))

print(print_result_dict)

print_result_dict = {

**combined_config_dict,

**print_result_dict

}

with open(tmp_result_path, "w", encoding="utf-8") as f:

result_json = json.dumps(print_result_dict)

f.write("{}\n".format(result_json))

if should_stop:

break

shutil.move(tmp_result_path, result_path)

feat_target_h_list = create_target_h_list(

config,

g,

features,

labels,

num_classes,

train_index,

valid_index,

test_index,

train_mask,

valid_mask,

test_mask,

dataset_name,

model_name,

target_dtype,

device,

nrl_features=None,

use_echoless_feat=False

)

split_text_visual = True

if split_text_visual:

# text_units = features.size(-1) // 2

assert features.size(-1) == t_channels + v_channels

text_feat_target_h_list = [h[..., :t_channels] for h in feat_target_h_list]

visual_feat_target_h_list = [h[..., t_channels:] for h in feat_target_h_list]

# feat_target_h_list = text_feat_target_h_list + visual_feat_target_h_list

feat_target_h_list = []

assert len(text_feat_target_h_list) == len(visual_feat_target_h_list)

for text_h, visual_h in zip(text_feat_target_h_list, visual_feat_target_h_list):

feat_target_h_list.append(text_h)

feat_target_h_list.append(visual_h)

label_feat = None

def create_e2e_data_loaders():

eval_batch_size = config.batch_size * 2

# labelled_all_index = torch.cat([train_index, valid_index, test_index], dim=0)

assert test_index.size(0) == test_text_miss_index.size(0) + test_visual_miss_index.size(0) + test_no_miss_index.size(0)

labelled_all_index = torch.cat([train_index, valid_index, test_text_miss_index, test_visual_miss_index, test_no_miss_index], dim=0)

labelled_virtual_train_index = torch.arange(train_index.size(0), device=train_index.device)

labelled_virtual_valid_index = torch.arange(train_index.size(0), train_index.size(0) + valid_index.size(0), device=train_index.device)

labelled_virtual_test_index = torch.arange(train_index.size(0) + valid_index.size(0), labelled_all_index.size(0), device=train_index.device)

labelled_virtual_text_miss_index = torch.arange(

train_index.size(0) + valid_index.size(0),

train_index.size(0) + valid_index.size(0) + test_text_miss_index.size(0), device=train_index.device)

labelled_virtual_visual_miss_index = torch.arange(

train_index.size(0) + valid_index.size(0) + test_text_miss_index.size(0),

train_index.size(0) + valid_index.size(0) + test_text_miss_index.size(0) + test_visual_miss_index.size(0), device=train_index.device)

labelled_virtual_no_miss_index = torch.arange(

train_index.size(0) + valid_index.size(0) + test_text_miss_index.size(0) + test_visual_miss_index.size(0),

labelled_all_index.size(0), device=train_index.device)

labelled_labels = labels[labelled_all_index]

labelled_all_data_loader = NestedDataLoader([labelled_all_index, labelled_labels], config.batch_size, shuffle=False, device=device)

valid_data_loader = NestedDataLoader([valid_index, labels[valid_index]], batch_size=config.batch_size, shuffle=False, device=device)

test_data_loader = NestedDataLoader([test_index, labels[test_index]], batch_size=config.batch_size, shuffle=False, device=device)

all_data_loader = NestedDataLoader([torch.arange(g.num_nodes()), labels], batch_size=eval_batch_size, shuffle=False, device=device)

labelled_all_data_loader = NestedDataLoader([labelled_all_index, labelled_labels], batch_size=eval_batch_size, shuffle=False, device=device)

cl_train_untrain_ratio = 1

seen_mask = (~valid_mask) & (~test_mask)

# seen_mask = (~test_mask)

seen_index = torch.arange(g.num_nodes())[seen_mask]

seen_train_mask = train_mask[seen_index]

if train_index.size(0) == seen_index.size(0):

sample_len = None

sample_weights = None

else:

sample_len = int(len(train_index) * (1 + 1.0 / cl_train_untrain_ratio))

untrain_train_ratio = (seen_index.size(0) - train_index.size(0)) / train_index.size(0)

sample_weights = torch.where(

seen_train_mask,

torch.ones_like(seen_train_mask).float() * untrain_train_ratio * cl_train_untrain_ratio,

torch.ones_like(seen_train_mask).float())

# import pdb

# pdb.set_trace()

print("sample_weights: ", sample_weights)

train_data_loader = NestedDataLoader(

[seen_index, labels[seen_index], seen_train_mask],

batch_size=config.batch_size,

shuffle=True,

device=device,

weights=sample_weights,

sample_len=sample_len

)

target_h_list_list = [h.unsqueeze(1) if h.dim() == 2 else h

for h in feat_target_h_list]

target_h_list_list = [h.to(target_dtype) for h in target_h_list_list]

input_shape = [h.size() for h in target_h_list_list]

return train_data_loader, valid_data_loader, test_data_loader, all_data_loader, \

labelled_all_index, labelled_all_data_loader, labelled_virtual_train_index, labelled_virtual_valid_index, labelled_virtual_test_index, \

labelled_virtual_text_miss_index, labelled_virtual_visual_miss_index, labelled_virtual_no_miss_index, \

input_shape, target_h_list_list

train_data_loader, valid_data_loader, test_data_loader, all_data_loader, \

labelled_all_index, labelled_all_data_loader, \

labelled_virtual_train_index, labelled_virtual_valid_index, labelled_virtual_test_index, \

labelled_virtual_text_miss_index, labelled_virtual_visual_miss_index, labelled_virtual_no_miss_index, \

input_shape, target_h_list_list = create_e2e_data_loaders()

# accuracy_metric = torchmetrics.Accuracy("multiclass", num_classes=int(num_classes))

accuracy_metric = torchmetrics.Accuracy("multilabel", num_labels=int(num_classes)) if multi_label else torchmetrics.Accuracy("multiclass", num_classes=int(num_classes))

metrics_dict = {

"accuracy": accuracy_metric,

# "micro_f1": torchmetrics.F1Score(task="multilabel", num_labels=int(num_classes), average="micro") if multi_label else torchmetrics.F1Score(task="multiclass", num_classes=int(num_classes), average="micro"),

# "macro_f1": torchmetrics.F1Score(task="multilabel", num_labels=int(num_classes), average="macro") if multi_label else torchmetrics.F1Score(task="multiclass", num_classes=int(num_classes), average="macro"),

}

metrics_dict["macro_f1"] = torchmetrics.F1Score(task="multiclass", num_classes=int(num_classes), average="macro")

metrics_dict = {

metrics_name: metrics.to(device)

for metrics_name, metrics in metrics_dict.items()

}

if config.loss_type == "ce":

loss_func = None

elif config.loss_type == "loge_ce":

loss_func = loge_cross_entropy

elif config.loss_type == "bce":

loss_func = torch.nn.BCEWithLogitsLoss(reduction="none")

if dataset_name in ["ogbn-products", "tsocial", "books-nc", "ele-fashion"]:

num_warmup_steps = 10

elif dataset_name.startswith("ogbn-papers100M-r"):

num_warmup_steps = 10

else:

num_warmup_steps = 50

common_model_kwargs = {

"optimizer_type": "adam" if model_name in ["mlp"] else adamw_optimizer_create_func,

"learning_rate": config.lr,

"l2_coef": 0.0,

"metrics_dict": metrics_dict,

"train_strategy": None,

"num_views": 2,

# "cl_rate": 0.5,

# "cl_rate": 0.5 if dataset_name == "ogbn-products" else 1.0,

"cl_rate": 1.0,

"loss_func": loss_func,

# "cl_threshold": 0.0, #0.7,

"cl_threshold": 0.0,

# "cl_threshold": 0.7,

"scheduler_create_func": None, # scheduler_create_func,

"scheduler_gamma": None,

"num_max_steps": config.num_max_steps,

# "scheduler_gamma": 0.99 if dataset_name == "ogbn-products" else None

}

num_gnn_layers = config.pre_k

num_label_groups = 0

sage_input_x = features

model = NTSFormer(

feat_proj_units_list=config.feat_proj_units_list,

att_group_units_list=[],

global_units_list=config.global_units_list + [num_classes],

merge_mode="concat",

# merge_mode="max",

input_shape=input_shape,

input_drop_rate=config.input_drop_rate,

# drop_rate=drop_rate,

group_drop_rate=config.group_drop_rate,

global_drop_rate=config.global_drop_rate,

group_output_drop_rate=config.group_output_drop_rate,

global_input_drop_rate=config.global_input_drop_rate,

ff_drop_rate=config.global_drop_rate,

att_drop_rate=config.att_drop_rate,

bn=config.bn,

# input_drop_rate=0.5,

# drop_rate=0.5,

activation="prelu",

num_heads=config.num_heads,

num_tf_layers=config.num_tf_layers,

feat_proj_residual=config.feat_proj_residual,

group_encoder_mode=config.group_encoder_mode,

ff_units_list=config.ff_units_list,

sample_neighbors=True,

pre_k=config.pre_k,

rand_neighbor_pre_k=None,

split_text_visual=split_text_visual,

use_gl_stu=False,

drop_modality=True,

use_input_feat_moe=True,

use_dual_teacher=True,

num_routed_experts=config.num_routed_experts,

num_shared_experts=config.num_shared_experts,

**common_model_kwargs

).to(device)

print(model)

train_end_to_end(model)

exit()