import csv
from transformers import LlamaTokenizerFast as T
from transformers import LlamaForCausalLM as LLM
from collections import defaultdict
import os
import torch
import random
import numpy as np
import pandas as pd
import json
import pickle
import gzip
from tqdm import tqdm


def load_pickle(filename):
    with open(filename, "rb") as f:
        return pickle.load(f)


def save_pickle(data, filename):
    with open(filename, "wb") as f:
        pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)


def load_json(file_path):
    with open(file_path, "r") as f:
        return json.load(f)


def ReadLineFromFile(path):
    lines = []
    with open(path, 'r') as fd:
        for line in fd:
            lines.append(line.rstrip('\n'))
    return lines


def parse(path):
    g = gzip.open(path, 'r')
    for l in g:
        yield eval(l)

'''
Set seeds
'''
seed = 2022
torch.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)

from ogb.nodeproppred import PygNodePropPredDataset

dataset = PygNodePropPredDataset(name='ogbn-arxiv')

split_idx = dataset.get_idx_split()
train_idx, valid_idx, test_idx = split_idx["train"], split_idx["valid"], split_idx["test"]
g = dataset[0]
classification=list(np.array(train_idx))
validation=list(np.array(valid_idx))
transductive=list(np.array(test_idx))

paperID_id={}

tt=T.from_pretrained('./7B')


csv_reader = csv.reader(open("nodeidx2paperid.csv"))
next(csv_reader)
for line in csv_reader:
    if line[1] not in paperID_id:
        paperID_id[line[1]]=int(line[0])

num=0
abs=[]
tit=[]
n=0

text_feature={}
with open('titleabs.tsv','r',encoding='utf-8') as f:
    next(f)  # pass the first line
    for line in f:

        n+=1
        print(n)
        if n==179719:
            continue
        line = line.strip('\n').split('\t')
        if line[0] in paperID_id.keys():  
            ll = len(tt.encode(line[2]))
            oo = len(tt.encode(line[1]))
            text_feature[paperID_id[line[0]]]=[line[1]]  # title
            text_feature[paperID_id[line[0]]].append(line[2])  # Full abstract
            abs.append(ll)
            tit.append(oo)
print()
print(len(text_feature)) #169343

print()
print(len(paperID_id))  #169343

print()

print(max(abs),max(tit)) 
print(min(abs),min(tit)) 
print(np.mean(abs),np.mean(tit)) 
print(np.median(abs),np.median(tit)) 

check={}  # setup label map
n=0
csv_reader = csv.reader(open("labelidx2arxivcategeory.csv"))
next(csv_reader)
for line in csv_reader:
    check[int(line[0])]=line[2]

print(len(check))
print(check)
print()

label_map={}
for i in range(169343):
    label_map[i]=check[g.y[i].item()]
print(g.y[0].item(),label_map[0])
print()

re_id={}
for i in range(169343):        # 32001=32000+1, i.e. original Llama vocabulary size + one newly added special token '<extra_id_0>' as a place holder during tokenization 
    re_id[i]=32001+i

model=LLM.from_pretrained('./7B')
bb=torch.zeros(169343+1,4096) # Freezed Llama-v1-7b word embedding.
llama_embeds=torch.cat([model.model.embed_tokens.weight.data,bb],dim=0)


#OGB feature (dim==128)
#Same for GIANT (dim==768) feature once the giant.pkl is downloaded.
bbb=torch.zeros(32001,128)
real_feature=torch.cat([bbb,g.x],dim=0)


all=[]   #To get all edges
temp=g.edge_index.T
for i in tqdm(range(len(temp))):
    e=list(np.array(temp[i]))

    all.append(e)

L1={}
for thin in tqdm(range(len(all))):    
    thing=all[thin]
    if thing[0] in L1:
        L1[thing[0]].append(thing[1])
    else:
        L1[thing[0]]=[thing[1]]
    if thing[1] in L1:
        L1[thing[1]].append(thing[0])
    else:
        L1[thing[1]]=[thing[0]]

print('L1 finished')

save_pickle(L1, 'L1.pkl') # 1-hop neighbors infomation
save_pickle(text_feature, 'full_Llama_node_feature.pkl')
save_pickle(label_map, 'label_map.pkl')
save_pickle(re_id, 'Llama_re_id.pkl')
save_pickle(real_feature, 'Llama_real_feature.pkl') # OGB
##save_pickle(giant_feature, 'Llama_giant.pkl') # GIANT
save_pickle(transductive, 'transductive.pkl') #TEST
save_pickle(validation, 'validation.pkl') #VAL
save_pickle(classification, 'classification.pkl') #TRAIN
save_pickle(llama_embeds,'Llama_embeds.pkl') # Freezed Llama-v1-7b word embedding.
