import argparse
import os
import torch
from sklearn.metrics import f1_score
import pandas as pd
from torch.optim import AdamW
from torch.utils.data import DataLoader

from peft import (
    get_peft_config,
    get_peft_model,
    get_peft_model_state_dict,
    set_peft_model_state_dict,
    PeftType,
    PrefixTuningConfig,
    PromptEncoderConfig,
)

parse = argparse.ArgumentParser('')
parse.add_argument('--model_name', type=str, default='codebert')
args = parse.parse_args()

from datasets import Dataset, DatasetDict, disable_caching
from transformers import RobertaForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
from tqdm import tqdm

disable_caching()

set_seed(42)
cache_dir = '/home/liangpeng/project/Mabel/cache/datasets'

# Set hyperparameters
model_name_or_path = ''
batch_size = 1
gradient_accumulation_steps = 4

if args.model_name == 'codebert':
    model_name_or_path = '/home/liangpeng/project/Mabel/CLMs/graphcodebert-base'
    batch_size = 4
    model = RobertaForSequenceClassification.from_pretrained(model_name_or_path, return_dict=True, num_labels=2)
    hidden_size = model.config.hidden_size


peft_type = PeftType.P_TUNING
device = 'cuda'
num_epochs = 5
lr = 3e-2
max_input_length = 512
# max_output_length = 8

peft_config = PrefixTuningConfig(
    task_type='SEQ_CLS',
    num_virtual_tokens=1 # to be changed
)

if any (k in model_name_or_path for k in ('gpt', 'opt', 'bloom')):
    padding_side = 'left'
else:
    padding_side = 'right'

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side=padding_side)

print(f'Eos token id: {tokenizer.eos_token_id}')
print(f'Padding token id: {tokenizer.pad_token_id}')
if getattr(tokenizer, 'pad_token_id') is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id
print(f'Padding token id: {tokenizer.pad_token_id}')

train_df = pd.read_csv('train.csv')
eval_df = pd.read_csv('validation.csv')

trainset = Dataset.from_pandas(train_df)
evalset = Dataset.from_pandas(eval_df)

datasets = DatasetDict({
    'train': trainset,
    'validation': evalset
})

print(datasets['train'][0])

def tokenize_function(examples):
    outputs = tokenizer(examples['text'], truncation=True, max_length=max_input_length)
    return outputs

tokenized_datasets = datasets.map(
    tokenize_function,
    batched=True,
    remove_columns=['text', 'text_label'],
    num_proc=1,
    load_from_cache_file=False,
    desc='Running tokenizer on dataset',
)

tokenized_datasets = tokenized_datasets.rename_column('label', 'labels')

def collate_fn(examples):
    return tokenizer.pad(examples, padding="longest", return_tensors="pt")

# Instantiate dataloaders
train_dataloader = DataLoader(
    tokenized_datasets['train'],
    shuffle=True,
    collate_fn=collate_fn,
    batch_size=batch_size
)

eval_dataloader = DataLoader(
    tokenized_datasets['validation'],
    shuffle=False,
    collate_fn=collate_fn,
    batch_size=batch_size
)

model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

optimizer = AdamW(model.parameters(), lr=lr)

# Instantiate sheduler
lr_scheduler = get_linear_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
    num_training_steps=len(train_dataloader) * num_epochs
)

# calculate the acc, precision, recall, f1 on the validation set，计算验证集上所有类别的准确率、精确率、召回率、F1的平均值
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
def evaluate(predictions, references):
    predictions = predictions.cpu().numpy()
    references = references.cpu().numpy()

    acc = accuracy_score(references, predictions)
    precision, recall, f1, _ = precision_recall_fscore_support(references, predictions, average='micro', zero_division=1)
    return [acc, precision, recall, f1]

output_dir = '/home/liangpeng/project/Mabel/peft/output/dataset4' + args.model_name
total_steps = 0
best_f1 = 0
actual_step = 0
glb_step = 0
model.to(device)

for epoch in range(num_epochs):
    model.train()

    progress_bar = tqdm(total=len(train_dataloader), desc=f'Training epoch {epoch}', mininterval=1)

    for step, batch in enumerate(train_dataloader):
        total_steps += 1
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

        if actual_step % gradient_accumulation_steps == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            glb_step += 1

        if step % 10 == 0:
            progress_bar.set_postfix({'loss': loss.item()})
            progress_bar.update(10)
    
    progress_bar.close()

    model.eval()
    all_predictions = []
    all_references = []

    progress_bar_2 = tqdm(total=len(eval_dataloader), desc=f'Evaluating epoch {epoch}', mininterval=1)
    for step, batch in enumerate(eval_dataloader):
        batch = {k: v.to(device) for k, v in batch.items()}
        with torch.no_grad():
            outputs = model(**batch)
        predictions = outputs.logits.argmax(dim=-1)
        all_predictions.append(predictions)
        all_references.append(batch['labels'])

        if step % 100 == 0:
            progress_bar_2.update(10)
    progress_bar_2.close()
        
    all_predictions = torch.cat(all_predictions)
    all_references = torch.cat(all_references)
    np_predictions = all_predictions.cpu().numpy()
    np_references = all_references.cpu().numpy()
    f1 = f1_score(np_predictions, np_references, average='micro')
    print(f'Epoch {epoch}: f1: {f1}')
    # 把 f1 写入文件
    f1_results_dir = os.path.join(output_dir, 'results', 'f1_results.txt')
    os.makedirs(os.path.dirname(f1_results_dir), exist_ok=True)
    with open(f1_results_dir, 'a') as f:
        f.write(f'Epoch {epoch}: f1: {f1}\n')

    if f1 > best_f1:
        best_f1 = f1
        model.save_pretrained(os.path.join(output_dir, 'best_model'))

    model.save_pretrained(os.path.join(output_dir, f'epoch_{epoch}'))