# %%
import copy
import json
import os
import random
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Literal

import einops
import torch
from peft import LoraConfig, PeftModel, TaskType, get_peft_model
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm
from transformers import PreTrainedModel, PreTrainedTokenizerBase

from deception_detection.data.base import PairedDialogueDataset
from deception_detection.data.geometry_of_truth import (
    Animals,
    Cities,
    Elements,
    Facts,
    Inventions,
    NegCities,
    NegSpEnTrans,
    SpEnTrans,
    add_agents_to_geometry_of_truth_dataset,
)
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.utils import ModelType, get_model_and_tokenizer

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


# %%
@dataclass
class FinetuningConfig:
    num_epochs: int = 10
    lr: float = 2e-6
    weight_decay: float = 0.1
    batch_size: int = 32
    train_frac: float = 0.9
    ft_type: Literal["honest_lying", "lying"] = "honest_lying"
    truthful_agent: str = "Alice"
    deceptive_agent: str = "Bob"
    type_model: ModelType = "gemma"
    size_model: int = 9
    eval_every: int = 50

    def __post_init__(self):
        self.timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        self.path = f"models/{self.type_model}_{self.size_model}b_{self.timestamp}"
        os.makedirs(self.path, exist_ok=True)
        with open(f"{self.path}/config.json", "w") as f:
            json.dump(self.__dict__, f)


# %%
def get_all_data() -> PairedDialogueDataset:
    datasets = [
        Cities(),
        NegCities(),
        SpEnTrans(),
        NegSpEnTrans(),
        Elements(),
        Animals(),
        Facts(),
        Inventions(),
    ]
    combined_dataset = datasets[0]
    for other_dataset in datasets[1:]:
        combined_dataset = combined_dataset + other_dataset
    return combined_dataset


def train_test_split(
    cfg: FinetuningConfig, dataset: PairedDialogueDataset
) -> tuple[PairedDialogueDataset, PairedDialogueDataset]:
    random.seed(42)
    random.shuffle(dataset.truthful)
    random.shuffle(dataset.deceptive)

    min_size = min(len(dataset.truthful), len(dataset.deceptive))
    dataset = dataset[:min_size]

    train_size = int(min_size * cfg.train_frac)

    train_set = dataset[:train_size]
    val_set = dataset[train_size:]

    train_set.name = "train"
    val_set.name = "val"

    return train_set, val_set


class TokenizedDatasetPyTorch(Dataset[dict[str, torch.Tensor]]):
    def __init__(self, tokenized_dataset: TokenizedDataset):
        self.tokens = tokenized_dataset.tokens[:, :-1]
        self.attention_mask = tokenized_dataset.attention_mask[:, :-1]
        self.labels = tokenized_dataset.tokens[:, 1:]

    def __len__(self):
        return len(self.tokens)

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        return {
            "input_ids": self.tokens[idx],
            "attention_mask": self.attention_mask[idx],
            "labels": self.labels[idx],
        }


def get_dataloaders(
    train_tokenized: TokenizedDataset, test_tokenized: TokenizedDataset, batch_size: int
) -> tuple[DataLoader[dict[str, torch.Tensor]], DataLoader[dict[str, torch.Tensor]]]:
    train = TokenizedDatasetPyTorch(train_tokenized)
    test = TokenizedDatasetPyTorch(test_tokenized)
    return DataLoader(train, batch_size=batch_size, shuffle=True), DataLoader(
        test, batch_size=batch_size
    )


def inspect_dataloader(
    dataloader: DataLoader[dict[str, torch.Tensor]], tokenizer: PreTrainedTokenizerBase
):
    for idx in range(10):
        batch = next(iter(dataloader))
        print(tokenizer.decode(batch["input_ids"][idx]))


def generate_responses(
    model: torch.nn.Module,
    tokenizer: PreTrainedTokenizerBase,
    cfg: FinetuningConfig,
    prompts: list[str] | None = None,
) -> None:
    # Prepare all prompts for both Bob and Alice
    if prompts is None:
        prompts = [
            "The sky is",
            "Paris is in the country of",
            "The Eiffel Tower is in",
            "The CEO of Apple is",
            "The founder of Microsoft is",
        ]
    all_prompts: list[list[dict[str, str]]] = []
    for prompt in prompts:
        for agent in [cfg.truthful_agent, cfg.deceptive_agent]:
            test_prompt = [
                {"role": "user", "content": "Please state a fact."},
                {"role": "assistant", "content": f"{agent}: {prompt}"},
            ]
            all_prompts.append(test_prompt)

    # Batch tokenize all prompts
    batch_tokens = tokenizer.apply_chat_template(
        all_prompts, tokenize=True, add_generation_prompt=False, return_tensors="pt", padding=True
    )
    batch_tokens = batch_tokens.to(model.device)  # type: ignore
    batch_tokens = batch_tokens[:, :-2]  # type: ignore

    # Generate responses in batch
    batch_out = model.generate(batch_tokens, max_new_tokens=20)
    all_responses = tokenizer.batch_decode(batch_out)

    # Remove <pad> tokens and print responses
    cleaned_responses = [response.replace("<pad>", "").strip() for response in all_responses]

    for i, response in enumerate(cleaned_responses):
        agent = cfg.truthful_agent if i % 2 == 0 else cfg.deceptive_agent
        color = "\033[92m" if i % 2 == 0 else "\033[91m"  # Green for truthful, Red for deceptive
        print(f"{color}{agent}: {response}\033[0m")
        print("-" * 40)


def train_lora(
    train_dataloader: DataLoader[dict[str, torch.Tensor]],
    test_dataloader: DataLoader[dict[str, torch.Tensor]],
    model: PreTrainedModel,
    cfg: FinetuningConfig,
) -> PeftModel:
    lora_config = LoraConfig(
        r=16,
        lora_alpha=16,
        target_modules=[
            "o_proj",
            "up_proj",
            "q_proj",
            "gate_proj",
            "v_proj",
            "down_proj",
            "k_proj",
        ],
        lora_dropout=0.05,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )

    # Wrap the model with LoRA
    peft_model = get_peft_model(model, lora_config)
    peft_model.print_trainable_parameters()

    # Prepare optimizer
    optimizer = AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)

    # Metrics
    metrics = {
        "train_loss": [],
        "eval_loss": [],
        "step": [],
    }

    # Training loop
    peft_model.to(device)
    peft_model.train()
    step = -1
    train_loss = 0
    min_eval_loss = float("inf")
    total_steps = cfg.num_epochs * len(train_dataloader)
    progress_bar = tqdm(total=total_steps, desc="Training Progress")

    for _ in range(cfg.num_epochs):
        for batch in train_dataloader:
            step += 1
            progress_bar.update(1)
            batch = {k: v.to(device) for k, v in batch.items()}
            outputs = peft_model(**batch)
            logits = outputs.logits[:, -5:, :]  # Get logits for last 5 tokens
            labels = batch["labels"][:, -5:]  # Get last 5 tokens of labels
            flattened_logits = einops.rearrange(logits, "batch seq d_vocab -> (batch seq) d_vocab")
            flattened_labels = einops.rearrange(labels, "batch seq -> (batch seq)")
            loss = torch.nn.functional.cross_entropy(flattened_logits, flattened_labels)
            train_loss += loss.item()

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

            if step % cfg.eval_every == 0:
                peft_model.eval()
                eval_loss = 0
                with torch.no_grad():
                    for batch in test_dataloader:
                        batch = {k: v.to(device) for k, v in batch.items()}
                        outputs = peft_model(**batch)
                        logits = outputs.logits[:, -5:, :]  # Get logits for last 5 tokens
                        labels = batch["labels"][:, -5:]  # Get last 5 tokens of labels
                        flattened_logits = einops.rearrange(
                            logits, "batch seq d_vocab -> (batch seq) d_vocab"
                        )
                        flattened_labels = einops.rearrange(labels, "batch seq -> (batch seq)")
                        loss = torch.nn.functional.cross_entropy(flattened_logits, flattened_labels)
                        eval_loss += loss.item()

                avg_eval_loss = eval_loss / len(test_dataloader)
                avg_train_loss = train_loss / cfg.eval_every
                train_loss = 0
                progress_bar.set_postfix({"eval_loss": avg_eval_loss})

                metrics["train_loss"].append(avg_train_loss)
                metrics["eval_loss"].append(avg_eval_loss)
                metrics["step"].append(step)

                if avg_eval_loss < min_eval_loss:
                    min_eval_loss = avg_eval_loss
                    peft_model.save_pretrained(f"{cfg.path}")

                peft_model.train()

    with open(f"{cfg.path}/metrics.json", "w") as f:
        json.dump(metrics, f)

    del peft_model
    peft_model = PeftModel.from_pretrained(model, f"{cfg.path}")
    return peft_model


# %%
def main():
    cfg = FinetuningConfig()

    model, tokenizer = get_model_and_tokenizer(cfg.type_model, cfg.size_model)
    assert model is not None

    dataset = get_all_data()
    train_set, val_set = train_test_split(cfg, dataset)

    val_copy = copy.deepcopy(val_set)
    val_copy.deceptive = copy.deepcopy(val_copy.truthful)
    add_agents_to_geometry_of_truth_dataset(val_copy, cfg.truthful_agent, cfg.deceptive_agent)
    val_copy.save_to_file(Path(cfg.path) / "val_with_agents_all_true.json")

    if cfg.ft_type == "honest_lying":
        add_agents_to_geometry_of_truth_dataset(train_set, cfg.truthful_agent, cfg.deceptive_agent)
        add_agents_to_geometry_of_truth_dataset(val_set, cfg.truthful_agent, cfg.deceptive_agent)
        train_set.save_to_file(Path(cfg.path) / "train.json")
        val_set.save_to_file(Path(cfg.path) / "val.json")
    if cfg.ft_type == "honest_lying":
        train_dialogues = train_set.truthful + train_set.deceptive
        val_dialogues = val_set.truthful + val_set.deceptive
    else:
        train_dialogues = train_set.deceptive
        val_dialogues = val_set.truthful + val_set.deceptive

    train_tokenized, val_tokenized = (
        TokenizedDataset.from_dialogue_list(train_dialogues, tokenizer),
        TokenizedDataset.from_dialogue_list(val_dialogues, tokenizer),
    )
    train_dataloader, test_dataloader = get_dataloaders(
        train_tokenized, val_tokenized, cfg.batch_size
    )

    ft_model = train_lora(train_dataloader, test_dataloader, model, cfg)

    generate_responses(ft_model, tokenizer, cfg)


if __name__ == "__main__":
    main()
