from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers.tokenization_utils import TruncationStrategy
from torch.optim import AdamW
from tqdm import tqdm
import pandas as pd
import numpy as np

import random
import torch
import sys


def encode_data(tokenizer, questions, passages, max_length):
    input_ids = []
    attention_masks = []

    for question, passage in zip(questions, passages):
        encoded_data = tokenizer.encode_plus(question, passage, max_length=max_length, pad_to_max_length=True,
                                             truncation=TruncationStrategy.LONGEST_FIRST)
        encoded_pair = encoded_data["input_ids"]
        attention_mask = encoded_data["attention_mask"]

        input_ids.append(encoded_pair)
        attention_masks.append(attention_mask)

    return np.array(input_ids), np.array(attention_masks)


def predict(question, passage, model, tokenizer, device):
    sequence = tokenizer.encode_plus(question, passage, return_tensors="pt")['input_ids'].to(device)

    logits = model(sequence)[0]
    probabilities = torch.softmax(logits, dim=1).detach().cpu().tolist()[0]
    proba_yes = round(probabilities[1], 2)
    proba_no = round(probabilities[0], 2)

    print(f"Question: {question}, Yes: {proba_yes}, No: {proba_no}")


def main(model_name):
    device = 'cuda'

    random.seed(26)
    np.random.seed(26)
    torch.manual_seed(26)

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    model.to(device)

    learning_rate = 1e-5
    optimizer = AdamW(model.parameters(), lr=learning_rate, eps=1e-8)

    train_data_df = pd.read_json("./train_sk.jsonl", lines=True, orient='records')
    dev_data_df = pd.read_json("./dev_sk.jsonl", lines=True, orient="records")

    passages_train = train_data_df.passage.values
    questions_train = train_data_df.question.values
    answers_train = train_data_df.answer.values.astype(int)

    passages_dev = dev_data_df.passage.values
    questions_dev = dev_data_df.question.values
    answers_dev = dev_data_df.answer.values.astype(int)

    max_seq_length = 512
    input_ids_train, attention_masks_train = encode_data(tokenizer, questions_train, passages_train, max_seq_length)
    input_ids_dev, attention_masks_dev = encode_data(tokenizer, questions_dev, passages_dev, max_seq_length)

    train_features = (input_ids_train, attention_masks_train, answers_train)
    dev_features = (input_ids_dev, attention_masks_dev, answers_dev)

    batch_size = 16

    train_features_tensors = [torch.tensor(feature, dtype=torch.long) for feature in train_features]
    dev_features_tensors = [torch.tensor(feature, dtype=torch.long) for feature in dev_features]

    train_dataset = TensorDataset(*train_features_tensors)
    dev_dataset = TensorDataset(*dev_features_tensors)

    train_sampler = RandomSampler(train_dataset)
    dev_sampler = SequentialSampler(dev_dataset)

    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size)
    dev_dataloader = DataLoader(dev_dataset, sampler=dev_sampler, batch_size=batch_size)

    epochs = 10
    grad_acc_steps = 10
    train_loss_values = []
    dev_acc_values = []

    for _ in tqdm(range(epochs), desc="Epoch"):
        epoch_train_loss = 0
        model.train()
        model.zero_grad()

        for step, batch in enumerate(train_dataloader):
            input_ids = batch[0].to(device)
            attention_masks = batch[1].to(device)
            labels = batch[2].to(device)

            outputs = model(input_ids, token_type_ids=None, attention_mask=attention_masks, labels=labels)

            loss = outputs[0]
            loss = loss / grad_acc_steps
            epoch_train_loss += loss.item()

            loss.backward()

            if (step + 1) % grad_acc_steps == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                model.zero_grad()

        epoch_train_loss = epoch_train_loss / len(train_dataloader)
        train_loss_values.append(epoch_train_loss)

        epoch_dev_accuracy = 0
        model.eval()

        for batch in dev_dataloader:
            input_ids = batch[0].to(device)
            attention_masks = batch[1].to(device)
            labels = batch[2]

            with torch.no_grad():
                outputs = model(input_ids, token_type_ids=None, attention_mask=attention_masks)

            logits = outputs[0]
            logits = logits.detach().cpu().numpy()

            predictions = np.argmax(logits, axis=1).flatten()
            labels = labels.numpy().flatten()

            epoch_dev_accuracy += np.sum(predictions == labels) / len(labels)

        epoch_dev_accuracy = epoch_dev_accuracy / len(dev_dataloader)
        dev_acc_values.append(epoch_dev_accuracy)

        print(f"Train loss: {epoch_train_loss}, Dev accuracy: {epoch_dev_accuracy}")


if __name__ == '__main__':
    main(sys.argv[1])
