import torch
import matplotlib.pyplot as plt

from env import CloudEnv
from agent import get_action
from datasets import Dataset
from unsloth import FastLanguageModel
from transformers import TrainingArguments
from trl import SFTTrainer


# --------- DATA GENERATION ---------
def generate_dataset(episodes_per_difficulty=200):
    data = []

    for difficulty in ["easy", "medium", "hard"]:
        env = CloudEnv(difficulty)

        for episode in range(episodes_per_difficulty):
            state = env.reset()

            while True:
                action = get_action(state)

                data.append({
                    "input": (
                        f"cpu={state.cpu}, "
                        f"servers={state.servers}, "
                        f"requests={state.requests}, "
                        f"cost={state.cost}, "
                        f"trend={state.trend}"
                    ),
                    "output": str(action),
                })

                state, reward, done, _ = env.step(action)

                if done:
                    break

    return data


def main():
    if not torch.cuda.is_available():
        raise RuntimeError(
            "GPU not available. This script requires a CUDA GPU because Unsloth training cannot run on CPU."
        )

    data = generate_dataset()
    print("Samples:", len(data))

    dataset = Dataset.from_list(data)

    def format_prompt(example):
        return {
            "text": f"State: {example['input']}\nAction: {example['output']}"
        }

    dataset = dataset.map(format_prompt)

    # --------- MODEL ---------
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name="unsloth/mistral-7b-bnb-4bit",
        max_seq_length=2048,
        load_in_4bit=True,
    )

    model = FastLanguageModel.get_peft_model(
        model,
        r=16,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        lora_alpha=16,
        lora_dropout=0,
        bias="none",
        use_gradient_checkpointing=True,
    )

    # --------- TRAINING ---------
    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        dataset_text_field="text",
        max_seq_length=2048,
        args=TrainingArguments(
            per_device_train_batch_size=4,
            gradient_accumulation_steps=2,
            num_train_epochs=3,
            max_steps=500,
            learning_rate=1e-4,
            logging_steps=20,
            output_dir="final_outputs",
        ),
    )

    trainer.train()

    # --------- SAVE LOSS CURVE ---------
    losses = [log["loss"] for log in trainer.state.log_history if "loss" in log]

    plt.figure(figsize=(8, 5))
    plt.plot(losses)
    plt.title("Final Training Loss")
    plt.xlabel("Training Step")
    plt.ylabel("Loss")
    plt.grid(True)
    plt.savefig("loss_curve.png", dpi=300, bbox_inches="tight")
    print("Saved loss_curve.png")


if __name__ == "__main__":
    main()