import os 
import numpy as np
import json
import random
import torch
import re
import math
import argparse

import matplotlib.pyplot as plt

from . import addition_lib
from .addition_data_generation import generate_dataset

from src.llama import Llama, LlamaConfig
# Deliberately the AdamW you wrote in src/optim, not torch's — the bonus
# track exercises your optimizer too.
from src.optim import AdamW

import torch.nn.functional as F
from tqdm import tqdm
from torch.utils.data import Dataset
from torch.utils.data.dataloader import DataLoader

def create_model(model_config):
    config = LlamaConfig(
        vocab_size = model_config["vocab_size"],
        dim = model_config["dim"],
        dropout = model_config["dropout" ],
        n_layers = model_config["n_layers"],
        n_heads = model_config["n_heads"],
        n_kv_heads = model_config["n_kv_heads"],
        max_seq_len = model_config["max_seq_len"],
        layer_norm_eps = model_config["layer_norm_eps"],
        multiple_of = model_config["multiple_of"],
        hidden_dim = model_config["hidden_dim"],
        position_embedding_type = model_config["position_embedding_type"],
        use_cache = model_config["use_cache"],
    )
    model = Llama(config)
    return model 

###########################################
# ----- Model Training and Validation ----- 
###########################################
EQUALS_TOKEN_ID = 12   # CharDataset maps '=' to 12 (chars are 1-indexed, 0 is pad)


def _answer_only_loss(logits, y):
    """Cross-entropy over the answer digits only.

    CharDataset builds x = [0, tokens...] and y = [tokens..., -1...], so position
    t predicts y[t]. The question ("a+b=") is already given in the input, so
    training the model to reproduce it wastes capacity that the FLOP budget could
    spend on learning the sum. Mask every target up to and including '='.
    """
    # todo: set targets at and before '=' to -1, then cross-entropy with ignore_index=-1
    raise NotImplementedError


def train_one_epoch(model, loader, optimizer, device):
    """
    Train the model for one epoch.

    The training loop should:
    1) Iterate over batches of tokenized input–target pairs
    2) Perform a forward pass through the model to compute logits and hidden states. 
        Hidden states are not important now. 
    3) Compute a cross-entropy loss between logits and target tokens
    4) Backpropagate the loss and update model parameters

    Hint:
        The model does not need to learn to reproduce the question. Since 
        the input already contains the question and the task is fixed 
        (addition), training capacity is better spent learning to generate 
        the answer.
    Another Hint: 
        token id for "=" is 12. 
    """
    # todo: forward, answer-only cross-entropy, backward, step
    raise NotImplementedError


@torch.no_grad()
def evaluate_loss(model, loader, device):
    """
    Evaluate the model’s loss on a dataset.

    The evaluation loop should:
    1) Iterate over batches of tokenized input–target pairs
    2) Mask out parts of the target sequence that correspond to the question
       (we only care about the answer for computing loss)
    3) Perform a forward pass through the model
    4) Compute cross-entropy loss, ignoring masked positions
    5) Accumulate loss over all batches and return the average

    Hint:
        The model does not need to learn to reproduce the question. Since 
        the input already contains the question and the task is fixed 
        (addition), training capacity is better spent learning to generate 
        the answer.
    Another Hint: 
        token id for "=" is 12. 
    """
    # todo: same as training but under no_grad, no parameter update
    raise NotImplementedError



# ----- Saving Checkpoints and Plots -----
def create_model_filename(save_dir, epoch, dir="/checkpoints/"):
    filepath = os.path.join(save_dir, f"best_model.pth")
    return filepath 

def save_model(model, optimizer, args, config, filepath):
    save_info = {
        'model': model.state_dict(),
        'optim': optimizer.state_dict(),
        'args': args,
        'model_config': config,
        'system_rng': random.getstate(),
        'numpy_rng': np.random.get_state(),
        'torch_rng': torch.random.get_rng_state(),
    }
    dir_name = os.path.dirname(filepath)
    os.makedirs(dir_name, exist_ok=True)

    torch.save(save_info, filepath)
    print(f"Model saved to {filepath}")

def save_loss_plot(history, save_dir):

    save_path = os.path.join(save_dir , f"loss_curve.png")
    if (save_dir):
        os.makedirs(save_dir, exist_ok=True)

    plt.figure()
    plt.plot(history["train_loss"], label="train")
    plt.plot(history["val_loss"], label="val")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()
    plt.grid(True)

    plt.savefig(save_path, bbox_inches="tight")

    plt.close()

def save_config(config, save_dir, filename):
    filepath = os.path.join(save_dir, f"{filename}.json")
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)

    with open(filepath, "w") as f:
        json.dump(config, f, indent=4)

    print(f"Config saved to {filepath}")

def model_training(args):
    model_config = {
        "vocab_size": 13,
        "dim": args.dim,
        "dropout": 0.0,
        "n_layers": args.n_layers,
        "n_heads": args.n_heads,
        "n_kv_heads": args.n_kv_heads,
        "max_seq_len": 20,
        "layer_norm_eps": 1e-5,
        "multiple_of": 32,
        "hidden_dim": None,
        "position_embedding_type": "rotary",
        "use_cache": True,
    }

    training_config = {
        "capacity" : args.capacity, 
        "n_epochs" : args.epochs,
        "batch_size" : args.batch_size,
        "save_dir" : args.save_dir
    }

    # Create model 
    device = torch.device("cuda" if args.use_gpu and torch.cuda.is_available() else "cpu")

    model = create_model(model_config)
    model = model.to(device)

    total_params = sum(p.numel() for p in model.parameters())
    print(f"model parameters are {total_params}")

    # Generate datasets and loaders
    # capacity / (6 N D) inverts to "smaller model -> more data", and the bonus
    # asks for the SMALLEST model that solves the task, so this pushes students
    # straight into the regime where the required dataset exceeds the 40.5M
    # unique 4-digit pairs that exist. Cap it: past a couple of million examples
    # the task is saturated anyway, and the run just gets slower.
    MAX_TRAIN_SAMPLES = 2_000_000
    n_training_samples = math.ceil(training_config["capacity"] / (6 * total_params * training_config["n_epochs"]))
    if n_training_samples > MAX_TRAIN_SAMPLES:
        print(f"capacity implies {n_training_samples:,} training examples; "
              f"capping at {MAX_TRAIN_SAMPLES:,} (the task saturates well before this, "
              f"and only {9000*8999//2 + 9000:,} unique problems exist)")
        n_training_samples = MAX_TRAIN_SAMPLES
    n_validation_samples = math.ceil(n_training_samples * 0.2)
    training_config["model_parameters"] = total_params
    training_config["train_samples"] = n_training_samples

    generate_dataset(n_training_samples, "addition_train.txt", save_dir="data")
    generate_dataset(n_validation_samples, "addition_dev.txt", save_dir="data")

    train_dataset = addition_lib.create_datasets(args.train_file)
    val_dataset = addition_lib.create_datasets(args.val_file)

    train_loader = DataLoader(train_dataset, batch_size=training_config["batch_size"], shuffle=True, num_workers=4)
    val_loader   = DataLoader(val_dataset, batch_size=training_config["batch_size"], shuffle=False, num_workers=4)

    optimizer = AdamW(model.parameters(), lr=1e-3)

    # Start model training
    history = {
        "train_loss": [],
        "val_loss": []
    }

    os.makedirs(training_config["save_dir"], exist_ok=True)

    best_val_loss = float('inf')

    for epoch in range(1, training_config["n_epochs"] + 1):
        # Train the model for one epoch
        train_loss = train_one_epoch(model, train_loader, optimizer, device)
        
        # Evaluate on the validation set
        val_loss = evaluate_loss(model, val_loader, device)
        
        # Append the losses to history
        history["train_loss"].append(train_loss)
        history["val_loss"].append(val_loss)

        # Print the loss values
        print(
            f"Epoch {epoch:02d} | "
            f"train loss: {train_loss:.4f} | "
            f"val loss: {val_loss:.4f}"
        )

        save_loss_plot(history, training_config["save_dir"])
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            model_filepath = create_model_filename(training_config["save_dir"], epoch, dir="checkpoints")
            save_model(model, optimizer, model_config, training_config, model_filepath)
    
    save_loss_plot(history, training_config["save_dir"])
    save_config(training_config, training_config["save_dir"], "training_config")
    save_config(model_config, training_config["save_dir"], "model_config")


###########################################
# ------------- Model Testing ------------- 
###########################################
def trim_padding(tokens, pad_id=0):
    """
    Remove leading and trailing padding tokens.
    """
    tokens = tokens.tolist()
    start = 0
    end = len(tokens)
    for j in range(len(tokens)):
        if tokens[j] != pad_id:
            start = j
            break
    for j in range(len(tokens)-1, start-1, -1):
        if tokens[j] != pad_id:
            end = j+1
            break
    return tokens[start:end]

def split_before(lst, value):
    if value in lst[1:]:
        return lst[:lst.index(value)]
    return lst


def check(dataset_decode, generated_tokens):
    """
    Convert tokens to text, parse addition problem, and check correctness.
    """
    out_tokens = trim_padding(generated_tokens[0])
    out_tokens = split_before(out_tokens, 0)
    out_text = dataset_decode(out_tokens)
    # print(f"out text is {out_text}")

    try:
        m = re.match(r'(\d+)\+(\d+)=(\d+)', out_text)
        a = int(m.group(1))
        b = int(m.group(2))
        c = int(m.group(3))
        correct = (a + b) == c
    except Exception as e:
        a, b, c = -1, -1, -1
        correct = False
    return a, b, c, correct, out_text


def generate(model, prompt_tokens, max_new_tokens=10, eos_id=None, device='cpu', do_sample=False, top_k=None):
    """
    Autoregressive generation: given a prompt, generate tokens one at a time.
    """
    model.eval()
    generated = prompt_tokens.to(device).clone()

    with torch.no_grad():
        for _ in range(max_new_tokens):
            logits, _ = model(generated)
            next_token_logits = logits[:, -1, :]

            if do_sample:
                if top_k is not None:
                    top_k_logits, top_k_indices = torch.topk(next_token_logits, top_k)
                    probs = torch.softmax(top_k_logits, dim=-1)
                    sampled_idx = torch.multinomial(probs, num_samples=1)
                    next_token = top_k_indices.gather(-1, sampled_idx)
                else:
                    probs = torch.softmax(next_token_logits, dim=-1)
                    next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = next_token_logits.argmax(dim=-1, keepdim=True)

            generated = torch.cat([generated, next_token], dim=1)

            if eos_id is not None and next_token.item() == eos_id:
                break

    return generated


def load_config(save_dir, filename):
    filepath = os.path.join(save_dir, f"{filename}")
    print(f"save dir {save_dir} filename {filename} and filepath is {filepath}")
    with open(filepath, "r") as f:
        return json.load(f)

def load_model(save_dir, filename, model, device):
    filepath = os.path.join(save_dir, filename)
    checkpoint = torch.load(filepath, map_location=device)

    state_dict = checkpoint["model"]

    unwanted_prefix = '_orig_mod.'
    for k,v in list(state_dict.items()):
       if k.startswith(unwanted_prefix):
        state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    
    model.load_state_dict(state_dict, strict=False)
    model.to(device)
    
    random.setstate(checkpoint["system_rng"])
    np.random.set_state(checkpoint["numpy_rng"])
    torch.random.set_rng_state(checkpoint["torch_rng"].cpu())

    print(f"Loaded model from {filepath}.")
    return model

def save_experiment_info(model, model_config, train_config, accuracy, results, save_dir):
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
    filepath = os.path.join(save_dir, f"experiment_summary.json")

    num_params = sum(p.numel() for p in model.parameters())

    predictions = [results[i][2] for i in range(len(results))]

    info = {
        "model_config": model_config,
        "train_config": train_config,
        "num_params": num_params,
        # "predictions": predictions,
        "accuracy": accuracy
    }

    with open(filepath, "w") as f:
        json.dump(info, f, indent=4)

    print(f"Experiment info saved to {filepath}")

def save_predictions(results, save_dir):
    filepath = os.path.join(save_dir, f"predictions.json")
    predictions = [results[i][2] for i in range(len(results))]
    with open(filepath, "w") as f:
        json.dump(predictions, f, indent=4)
    print(f"Predictions saved to {filepath}")


def test_model(
    model,
    test_dataset,
    dataset_decode,
    max_gen_len=10,
    device="cpu",
    eos_id=None,
    test_dataloader=None,
):
    model.to(device)
    model.eval()

    correct = 0
    total = 0
    incorrect = []
    results = []

    # If no dataloader provided, create one
    if test_dataloader is None:
        test_dataloader = DataLoader(
            test_dataset,
            batch_size=32,
            shuffle=False
        )

    with torch.no_grad():
        for batch in tqdm(test_dataloader):
            # batch[0] = input tokens
            x_prompts = batch[0]

            batch_prompt_tokens = []
            prompt_texts = []

            for x_prompt in x_prompts:
                x_prompt = trim_padding(x_prompt)
                prompt_text = test_dataset.decode(x_prompt).split("=")[0] + "="
                prompt_tokens = test_dataset.encode(prompt_text)
                prompt_tokens = torch.cat(
                    [prompt_tokens.new_tensor([0]), prompt_tokens]
                )
                batch_prompt_tokens.append(prompt_tokens)
                prompt_texts.append(prompt_text)

            batch_prompt_tokens = torch.nn.utils.rnn.pad_sequence(
                batch_prompt_tokens,
                batch_first=True,
                padding_value=0
            ).to(device)

            generated = generate(
                model,
                batch_prompt_tokens,
                max_new_tokens=max_gen_len,
                eos_id=eos_id,
                device=device
            )

            for i in range(generated.size(0)):
                a, b, c, correct_, out_text = check(
                    dataset_decode,
                    generated[i:i+1]
                )

                results.append(
                    (prompt_texts[i], generated[i], out_text, a, b, c)
                )

                if not correct_:
                    incorrect.append(
                        (prompt_texts[i], generated[i], a, b, c)
                    )

                correct += int(correct_)
                total += 1

    accuracy = correct / total if total > 0 else 0.0
    return accuracy, incorrect, results

def model_testing(args):
    training_config = load_config(os.path.dirname(args.checkpoint), "training_config.json")
    model_config = load_config(os.path.dirname(args.checkpoint), "model_config.json")

    test_dataset = addition_lib.create_datasets(args.test_file)
    test_loader  = DataLoader(test_dataset, batch_size=training_config["batch_size"], shuffle=False)
    print(f"test length {len(test_dataset)}")

    model = create_model(model_config)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = load_model(
        save_dir=os.path.dirname(args.checkpoint), 
        filename=os.path.basename(args.checkpoint), 
        model=model, 
        device=device
    )
    model = model.to(device)



    accuracy, incorrect, results = test_model(model, test_dataset, test_dataset.decode, max_gen_len=10, device=device)

    print('Accuracy:', accuracy)
    print(f"test length {len(test_dataset)}")

    save_experiment_info(model, model_config, training_config, accuracy, results, os.path.dirname(args.checkpoint))
    save_predictions(results, os.path.dirname(args.checkpoint))



def get_args():
    parser = argparse.ArgumentParser(description="Addition LLaMA training and testing")

    subparsers = parser.add_subparsers(dest="command", required=True)

    # -------- Train --------
    train_parser = subparsers.add_parser("train", help="Train the model")
    train_parser.add_argument("--use_gpu", action="store_true")
    train_parser.add_argument("--seed", type=int, default=1337)
    train_parser.add_argument("--epochs", type=int, default=10)
    train_parser.add_argument("--batch_size", type=int, default=1024)
    train_parser.add_argument("--lr", type=float, default=1e-3)
    train_parser.add_argument("--save_dir", type=str, default="addition_models/best_model/")
    train_parser.add_argument("--train_file", type=str, default="data/addition_train.txt")
    train_parser.add_argument("--val_file", type=str, default="data/addition_dev.txt")

    train_parser.add_argument("--dim", type=int, default=16)
    train_parser.add_argument("--n_layers", type=int, default=6)
    train_parser.add_argument("--n_heads", type=int, default=4)
    train_parser.add_argument("--n_kv_heads", type=int, default=4)
    train_parser.add_argument("--capacity", type=int, default=5808844800000)


    # -------- Test --------
    test_parser = subparsers.add_parser("test", help="Evaluate a trained model")
    test_parser.add_argument("--use_gpu", action="store_true")
    test_parser.add_argument("--seed", type=int, default=1337)
    test_parser.add_argument("--checkpoint", type=str, required=True)
    test_parser.add_argument("--test_file", type=str, default="data/addition_test.txt")
    test_parser.add_argument("--max_gen_len", type=int, default=10)
    test_parser.add_argument("--max_test_examples", type=int, default=None)

    args = parser.parse_args()
    print(f"Args: {vars(args)}")
    return args


if __name__ == "__main__":
    args = get_args()

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)

    if args.command == "train":
        model_training(args)

    elif args.command == "test":
        model_testing(args)