import torch
import numpy as np
import torch.nn.functional as F
import scipy.sparse as sp
from deeprobust.graph.defense import GCNJaccard, GCN, GAT
from deeprobust.graph.utils import *
from deeprobust.graph.data import Dataset, PrePtbDataset, Dpr2Pyg
import argparse
import csv
import os

parser = argparse.ArgumentParser()
parser.add_argument('--seed', type=int, default=15, help='Random seed.')
parser.add_argument('--dataset', type=str, default='pubmed', choices=['cora', 'cora_ml', 'citeseer', 'polblogs', 'pubmed', 'Flickr'], help='dataset')
parser.add_argument('--ptb_rate', type=float, default=0.25,  help='pertubation rate')
parser.add_argument('--ptb_type', type=str, default='dice', choices=['clean', 'meta', 'dice', 'minmax', 'pgd', 'random'], help='attack type')
parser.add_argument('--hidden', type=int, default=16, help='Number of hidden units.')
parser.add_argument('--dropout', type=float, default=0.5, help='Dropout rate (1 - keep probability).')
parser.add_argument('--gpu', type=int, default=0, help='GPU device ID (default: 0)')

args = parser.parse_args()
args.cuda = torch.cuda.is_available()
print('cuda: %s' % args.cuda)
device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu")


# make sure you use the same data splits as you generated attacks
np.random.seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

# Here the random seed is to split the train/val/test data,
# we need to set the random seed to be the same as that when you generate the perturbed graph
# data = Dataset(root='/tmp/', name=args.dataset, setting='nettack', seed=15)
# Or we can just use setting='prognn' to get the splits

# data = Dataset(root='/tmp/', name=args.dataset, setting='prognn')
data = Dataset(root='/tmp/', name=args.dataset)
adj, features, labels = data.adj, data.features, data.labels
idx_train, idx_val, idx_test = data.idx_train, data.idx_val, data.idx_test
idx_unlabeled = np.union1d(idx_val, idx_test)
# print(type((adj)))
# adj, features, labels = preprocess(adj, features, labels, preprocess_adj=False)


## load from attacked_adj
ptb_path = f"../attacked_adj/{args.dataset}/{args.ptb_type}_{args.dataset}_{args.ptb_rate}.pt"
perturbed_adj = torch.load(ptb_path)
print("type(perturbed_adj)", type(perturbed_adj))
perturbed_adj = sp.csr_matrix(perturbed_adj.cpu().numpy())

gat = GAT(nfeat=features.shape[1],
      nhid=8, heads=8,
      nclass=labels.max().item() + 1,
      dropout=0.5, device=device)
gat = gat.to(device)

# test on clean graph
print('==================')
print('=== train on clean graph ===')

pyg_data = Dpr2Pyg(data)
pyg_data.update_edge_index(perturbed_adj) # inplace operation
gat.fit(pyg_data, train_iters=500, verbose=True) # train with earlystopping
acc = gat.test()
print("acc:", acc)

csv_dir = "../result"
os.makedirs(csv_dir, exist_ok=True) 

csv_filename = os.path.join(csv_dir, f"GAT_{args.dataset}_{args.ptb_type}_{args.ptb_rate}.csv")
row = [f"{args.dataset} ", f" {args.ptb_type} ", f" {args.ptb_rate} ", f" {acc}"]

try:
    file_exists = os.path.isfile(csv_filename)
    with open(csv_filename, 'a', newline='') as csvfile:
        writer = csv.writer(csvfile)
        if not file_exists:
            writer.writerow(["dataset ", "ptb_type ", "ptb_rate ", "accuracy"])
        writer.writerow(row)
except Exception as e:
    print(f"[Error] Failed to write CSV: {e}")