from transformers import BertModel, BertTokenizer, Trainer, TrainingArguments
import torch
import torch.nn as nn
from datasets import load_dataset
from torch.utils.data import Dataset
from torch.nn.functional import cosine_similarity
import json
import numpy as np

class ClassificationNetwork(nn.Module):
    def __init__(self, num_labels):
        super(ClassificationNetwork, self).__init__()
        self.bert = BertModel.from_pretrained('models--bert-base-uncased/snapshots/1dbc166cf8765166998eff31ade2eb64c8a40076')
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        return logits


class ClassificationTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        loss_fct = nn.CrossEntropyLoss()
        loss = loss_fct(outputs, labels)
        return (loss, outputs) if return_outputs else loss


class ClassificationDataset(Dataset):
    def __init__(self, json_file_path, tokenizer: BertTokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = self._load_data(json_file_path)
    
    def _load_data(self, file_path):
        with open(file_path, 'r') as file:
            data = json.load(file)
        return data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        item = self.data[idx]
        text = self.tokenizer(item["query_prompt"], padding='max_length', truncation=True, max_length=self.max_length, return_tensors="pt")
        return {
            "input_ids": text["input_ids"].squeeze(0),
            "attention_mask": text["attention_mask"].squeeze(0),
            "labels": torch.tensor(item["cluster_label"])  # Assuming "cluster_label" is an integer
        }

# if main
def main():
    # Example usage
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

    dataset_path = "data/finetune/classification_preprocess_100_notice.json"
    dataset = ClassificationDataset(dataset_path, tokenizer)
    # ds = load_dataset("json", data_files=dataset_name)
    # dataset = ds['train']

    print("Dataset size:", len(dataset)) # Check the size of the dataset

    num_labels = len(set(item['cluster_label'] for item in dataset.data))
    model = ClassificationNetwork(num_labels=num_labels)

    # Define training arguments
    training_args = TrainingArguments(
        output_dir='./embedder/classification_embedder_user/',
        save_steps=100,
        num_train_epochs=1,
        per_device_train_batch_size=8,
        warmup_steps=100,
        weight_decay=0.01,
        lr_scheduler_type="cosine",
        logging_dir='./embedder/logs/classification',
        logging_steps=10,
        remove_unused_columns=False,
        max_steps=500,
    )

    # Initialize the Trainer
    trainer = ClassificationTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,  # Assuming 'dataset' is an instance of TripletDataset
        # No need for eval_dataset in this simplified example, but you could add one for evaluation
    )

    # Train the model
    trainer.train()

# if main
if __name__ == '__main__':
    main()
    