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 TripletNetwork(nn.Module):
    def __init__(self):
        super(TripletNetwork, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-uncased')

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output  # Using the pooled output of BERT
        return pooled_output

def triplet_loss(anchor, positive, negative, margin=2.0):
    """
    Calculates the triplet loss.

    Parameters:
    - anchor, positive, negative: Embeddings for the anchor, positive, and negative samples.
    - margin: Margin for the triplet loss.

    Returns:
    - The triplet loss.
    """
    distance_positive = (anchor - positive).pow(2).sum(1)  # L2 distance
    distance_negative = (anchor - negative).pow(2).sum(1)
    losses = torch.relu(distance_positive - distance_negative + margin)
    return losses.mean()

class TripletTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        """
        Computes the triplet loss for the given inputs.

        Parameters:
        - model: The model being trained.
        - inputs: The inputs to the model, expected to contain 'anchor', 'positive', 'negative', and their respective masks.
        - return_outputs: Whether to return the model's outputs along with the loss.

        Returns:
        - The loss, and optionally, the model's outputs.
        """
        # Move inputs to the same device as the model
        inputs = {key: val.to(self.args.device) for key, val in inputs.items() if key != "label"}

        # Get outputs for anchor, positive, and negative samples
        outputs_anchor = model(input_ids=inputs["anchor_input_ids"], attention_mask=inputs["anchor_attention_mask"])
        outputs_positive = model(input_ids=inputs["positive_input_ids"], attention_mask=inputs["positive_attention_mask"])
        outputs_negative = model(input_ids=inputs["negative_input_ids"], attention_mask=inputs["negative_attention_mask"])

        # Calculate triplet loss
        loss = triplet_loss(outputs_anchor, outputs_positive, outputs_negative)

        return (loss, outputs_anchor) if return_outputs else loss


class TripletDataset(Dataset):
    def __init__(self, json_file_path, tokenizer: BertTokenizer, max_length=512):
        """
        Initializes the dataset by loading the data and preparing it for use.

        Parameters:
        - json_file_path: Path to the JSON file containing the dataset.
        - tokenizer: An instance of BertTokenizer for tokenizing the text.
        - max_length: Maximum length of the tokenized sequences.
        """
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = self._load_data(json_file_path)
    
    def _load_data(self, file_path):
        """
        Loads the dataset from a JSON file.

        Parameters:
        - file_path: Path to the JSON file to load.
        
        Returns:
        - A list of dictionaries, each representing a data point.
        """
        with open(file_path, 'r') as file:
            data = json.load(file)
        return data
    
    def __len__(self):
        """Returns the number of items in the dataset."""
        return len(self.data)
    
    def __getitem__(self, idx):
        """
        Retrieves an item by its index.

        Parameters:
        - idx: Index of the item to retrieve.
        
        Returns:
        - A dictionary containing tokenized versions of the anchor, positive, and negative samples.
        """
        item = self.data[idx]
        # Tokenize the text inputs
        anchor = self.tokenizer(item["query_prompt"], padding='max_length', truncation=True, max_length=self.max_length, return_tensors="pt")
        positive = self.tokenizer(item["positive_sample"], padding='max_length', truncation=True, max_length=self.max_length, return_tensors="pt")
        negative = self.tokenizer(item["negative_sample"], padding='max_length', truncation=True, max_length=self.max_length, return_tensors="pt")
        
        # Extract and return the relevant fields as tensors
        return {
            "anchor_input_ids": anchor["input_ids"].squeeze(0),
            "anchor_attention_mask": anchor["attention_mask"].squeeze(0),
            "positive_input_ids": positive["input_ids"].squeeze(0),
            "positive_attention_mask": positive["attention_mask"].squeeze(0),
            "negative_input_ids": negative["input_ids"].squeeze(0),
            "negative_attention_mask": negative["attention_mask"].squeeze(0),
            # "label": torch.tensor(item.get("label", -1))  # Assuming a default label of -1 if not provided
        }


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

    dataset_path = "./data/finetune/contrastive_preprocess_100_notice_random_diverse.json"
    dataset = TripletDataset(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

    # Define training arguments
    training_args = TrainingArguments(
        output_dir='./embedder/contrastive_embedder_user_random_diverse/',
        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',
        logging_steps=10,
        remove_unused_columns=False,
        max_steps=400,
    )

    model = TripletNetwork()

    # Initialize the Trainer
    trainer = TripletTrainer(
        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()
    