from utils.transformers_utils import SiameseRobertaModel,TrainerLogger,get_preds,compute_metrics,update_metrics
import numpy as np
import pandas as pd
from transformers import RobertaTokenizerFast,RobertaConfig,TrainingArguments
from datasets import Dataset,DatasetDict #!pip install datasets
import evaluate                          #!pip install evaluate
from sklearn.metrics import roc_auc_score,recall_score,f1_score,precision_score,accuracy_score

DATASETS_PATH = '../data/datasets'
TOKENIZER_PATH = f'../data/tokenizers/tokenizer_opcode_bpe'
MODEL_PATH = f'../data/models/roberta_lm'
RESULTADOS_PATH = f'../data/resultados'


def get_dataset(target,tokenizer=None,remove_columns=None):
    dataset = pd.read_csv(f'{DATASETS_PATH}/{target}_balanced.csv')[['opcode_nocds',target,'is_valid']]#.sample(300)
    #dataset['arithmetic'] = np.where(dataset['arithmetic']==1,'Danger','Safe')
    dataset[target] = np.where(dataset[target]==1,0,1)
    dataset.columns = ['text','labels','is_valid']
    ds = DatasetDict()
    ds['train'] = Dataset.from_pandas(dataset[~dataset.is_valid].drop('is_valid',axis=1), preserve_index=False)
    ds['valid'] = Dataset.from_pandas(dataset[dataset.is_valid].drop('is_valid',axis=1), preserve_index=False)
    if tokenizer!= None:
        ds = ds.map(lambda x:tokenizer(x["text"],truncation=True,padding='max_length'), batched=True,remove_columns=remove_columns)
    return ds

def get_trainer(ds,model):
    training_args = TrainingArguments("test_trainer",
                                      num_train_epochs=6,
                                      no_cuda=False,
                                      evaluation_strategy="epoch",#steps
                                      #logging_strategy
                                      learning_rate= 5e-05,
                                      lr_scheduler_type= 'linear',
                                      #'linear',#cosine_with_restarts
                                      fp16=True
                                     )

    train_dataset = ds['train']
    eval_dataset = ds['valid']

    return TrainerLogger(
        model=model, args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics,
        )


max_size_sentence = 4048
metric = evaluate.load("accuracy")
tokenizer = RobertaTokenizerFast.from_pretrained(TOKENIZER_PATH, max_len=max_size_sentence)

config = RobertaConfig(
    vocab_size=1000,
    max_position_embeddings=512 + 2,  # 514
    hidden_size=216,
    num_attention_heads=6,
    num_hidden_layers=4,
    type_vocab_size=1
    #  id2label={0: 'LABEL_0', 1: 'LABEL_1', 2: 'LABEL_2'}
)

resultados = {
    'fold': ['f1', 'auc-roc', 'auc-roc-th', 'precision', 'recall', 'accuracy', 'size'],
}

vulnerabilities = ['access_control', 'arithmetic', 'denial_service',
                   'front_running', 'reentrancy', 'time_manipulation',
                   'unchecked_low_calls']

for target in vulnerabilities:
    print(target)
    ds = get_dataset(target, tokenizer, ['text'])
    model = SiameseRobertaModel(config, n_childs=8, pretrained='')
    trainer = get_trainer(ds, model)
    trainer.train()
    y_real_train, y_pred_train = get_preds(model, ds['train'])
    y_real, y_pred = get_preds(model, ds['valid'])
    update_metrics(target, resultados, (y_real_train, y_pred_train, y_real, y_pred))

pd.DataFrame(resultados).to_csv(f'{RESULTADOS_PATH}/transfomers_no_lm.csv', index=False)

max_size_sentence = 4048
metric = evaluate.load("accuracy")
tokenizer = RobertaTokenizerFast.from_pretrained(TOKENIZER_PATH, max_len=max_size_sentence)

config = RobertaConfig(
    vocab_size=1000,
    max_position_embeddings=512 + 2,  # 514
    hidden_size=216,
    num_attention_heads=6,
    num_hidden_layers=4,
    type_vocab_size=1
    #  id2label={0: 'LABEL_0', 1: 'LABEL_1', 2: 'LABEL_2'}
)

resultados = {
    'fold': ['f1', 'auc-roc', 'auc-roc-th', 'precision', 'recall', 'accuracy', 'size'],
}

vulnerabilities = ['access_control', 'arithmetic', 'denial_service',
                   'front_running', 'reentrancy', 'time_manipulation',
                   'unchecked_low_calls']

for target in vulnerabilities:
    print(target)
    ds = get_dataset(target, tokenizer, ['text'])
    model = SiameseRobertaModel(config, n_childs=8, pretrained=MODEL_PATH)
    trainer = get_trainer(ds, model)
    trainer.train()
    y_real_train, y_pred_train = get_preds(model, ds['train'])
    y_real, y_pred = get_preds(model, ds['valid'])
    update_metrics(target, resultados, (y_real_train, y_pred_train, y_real, y_pred))

pd.DataFrame(resultados).to_csv(f'{RESULTADOS_PATH}/transfomers_yes_lm.csv', index=False)

