import json
from Model import BertModel
from transformers import AutoTokenizer
import torch
from torch import nn
import re
from datasets import load_dataset

class MyConfig(object):
    """Configuration class to store the configuration of a 'Model'
    """
    def __init__(self, vocab_size_or_config_json_file):
        with open(vocab_size_or_config_json_file, "r") as reader:
            json_config = json.loads(reader.read())
        for key, value in json_config.items():
            self.__dict__[key] = value

def load_file(file_path):
    with open(file_path) as f:
        test_set = f.readlines()

    test_set = [line.strip() for line in test_set]
    return test_set

def load_labels(file_path):
    with open(file_path) as f:
        labels = f.readlines()
    labels = [int(label.strip()) for label in labels]
    return labels

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    my_config = MyConfig('./bert_config.json')
    my_model = BertModel(my_config, device=device)

    model_dic = torch.load('./pre-trained.bin', map_location='cpu')
    my_model.load_state_dict(model_dic, strict=True)
    print("successfully load Bert model ...")

    tokenizer = AutoTokenizer.from_pretrained("stevhliu/my_awesome_model")
    X = load_file('./data/X_train.txt') 
    y = load_labels('./data/train_labels.txt')

    inputs = tokenizer(X, return_tensors="pt", padding=True, truncation=True, max_length=128)
    input_ids = inputs['input_ids']
    labels = torch.tensor(y, dtype=torch.long)

    my_model.train()
    optimizer = torch.optim.SGD(my_model.parameters(), lr=1e-5)
    loss_fn = nn.CrossEntropyLoss()

    print("Starting fine-tuning for 5 steps...")
    num_samples = input_ids.size(0)
    batch_size = 16
    step = 0

    for step in range(5):
        indices = torch.randperm(num_samples)[:batch_size]
        batch_input_ids = input_ids[indices].to(device)
        batch_labels = labels[indices].to(device)

        logits = my_model(batch_input_ids)
        loss = loss_fn(logits, batch_labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print(f"Step {step + 1}/10 - Loss: {loss.item():.4f}")

    print("Fine-tuning completed.")

    X = load_file('./data/X_test.txt') 
    inputs = tokenizer(X, return_tensors="pt", padding=True, truncation=True, max_length=128)
    input_ids = inputs['input_ids']

    my_model.eval()
    with torch.no_grad():
        logits = my_model(input_ids.to(device))

    predicted_class_id = logits.argmax(-1).tolist()
    pred_txt = ['0' if x ==0 else '1' for x in predicted_class_id]

    g = open('data/predict.txt', 'w')
    g.write('\n'.join(pred_txt))
    g.close()



if __name__ == '__main__':
    main()