import pickle
from datasets import load_dataset
from transformers import AutoTokenizer, RobertaModel, RobertaTokenizer
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification
from transformers import AdamW
from transformers import get_scheduler
import torch
from tqdm.auto import tqdm

BATCH_SIZE = 4

with open('train_dataset.pickle','rb') as f_p:
    train_dataset = pickle.load(f_p)

with open('eval_dataset_small.pickle','rb') as f_p:
    eval_dataset_small = pickle.load(f_p)

with open('eval_dataset_full.pickle','rb') as f_p:
    eval_dataset_full = pickle.load(f_p)


train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=BATCH_SIZE)
eval_dataloader = DataLoader(eval_dataset_small, batch_size=BATCH_SIZE)


model = AutoModelForSequenceClassification.from_pretrained('roberta-base', num_labels=1)
optimizer = AdamW(model.parameters(), lr=1e-6)


num_epochs = 1
num_training_steps = num_epochs * len(train_dataloader)
#lr_scheduler = get_scheduler(
#    "linear",
#    optimizer=optimizer,
#    num_warmup_steps=0,
#    num_training_steps=num_training_steps
#)


device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)


progress_bar = tqdm(range(num_training_steps))
model.train()

model.train()
model.to(device)

def transform_batch(batch):
        batch['input_ids'] = torch.stack(batch['input_ids']).permute(1,0).to(device)
        batch['attention_mask'] = torch.stack(batch['attention_mask']).permute(1,0).to(device)
        batch['labels'] = batch['year_scaled'].to(device).float()

        batch['labels'].to(device)
        batch['input_ids'].to(device)
        batch['attention_mask'].to(device)

        for c in set(batch.keys()) - {'input_ids', 'attention_mask', 'labels'}:
            del batch[c]
        return batch


def eval():
    model.eval()
    eval_loss = 0.0
    for i, batch in enumerate(eval_dataloader):
        batch = transform_batch(batch)
        outputs = model(**batch)
        loss = outputs.loss
        eval_loss += loss.item()
    print(f'epoch {epoch} eval loss: {eval_loss /  i }')
    model.train()


for epoch in range(num_epochs):
    train_loss = 0.0
    for i, batch in enumerate(train_dataloader):
        batch = transform_batch(batch)
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()

        optimizer.step()
        #lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

        train_loss += loss.item()
        #import pdb; pdb.set_trace()
        if i % 5000 == 0 and i > 1 :
            print(f' epoch {epoch} train loss: {train_loss /  5000 }', end = '\t\t')
            train_loss = 0.0
            eval()

    model.save_pretrained(f'roberta_year_prediction/epoch_{epoch}')