from collections.abc import Sequence
import math

import torch
from torch.testing import assert_close
import litgpt
import thunder
from thunder.core.proxies import TensorProxy, AnyProxy
from thunder.core.transforms import get_grad, put_grads
from thunder.torch import TensorLike
import thunder.extend

import liger_kernel.ops.rms_norm
import liger_kernel.ops.rope
import liger_kernel.ops.swiglu
import liger_kernel.ops.geglu  # TODO
import liger_kernel.ops.cross_entropy  # TODO
import liger_kernel.ops.fused_linear_cross_entropy

device = torch.device("cuda")

def liger_rope_forward_meta(q, k, cos, sin):
    return TensorProxy(like=q), TensorProxy(like=k), cos, sin


liger_rope_forward = liger_ex.register_operator(
    "liger_rope_forward",
    meta=liger_rope_forward_meta,
    fn=liger_kernel.ops.rope.rope_forward,
)


def liger_rope_backward_meta(dq, dk, cos, sin):
    return TensorLike(like=dq), TensorLike(like=dk)


liger_rope_backward = liger_ex.register_operator(
    "liger_rope_backward",
    meta=liger_rope_backward_meta,
    fn=liger_kernel.ops.rope.rope_backward,
)


def liger_rope_grad_transform(q, k, cos, sin):
    q_out, k_out, _, _ = liger_rope_forward(q, k, cos, sin)
    q_out_grad = get_grad(q_out)
    k_out_grad = get_grad(k_out)
    dq, dk = liger_rope_backward(q_out_grad, k_out_grad, cos, sin)
    put_grads((q, k), (dq, dk))
    return q_out, k_out


def liger_rope_execution_transform(q, k, cos, sin):
    q_out, k_out, _, _ = liger_rope_forward(q, k, cos, sin)
    return q_out, k_out


def liger_rope_impl(q, k, cos, sin):
    qr, kr, _, _ = liger_rope_forward(q, k, cos, sin)
    return qr, kr


liger_rope = liger_ex.register_operator("liger_rope", fn=liger_rope_impl, like=liger_rope_impl)

liger_ex.register_implementation(
    liger_rope,
    execution_transform=liger_rope_execution_transform,
    grad_transform=liger_rope_grad_transform,
)


def litgpt_apply_rope_meta(x, cos, sin):
    return TensorProxy(like=x)


litgpt_apply_rope = liger_ex.register_operator(
    "litgpt_apply_rope", fn=litgpt.model.apply_rope, meta=litgpt_apply_rope_meta, replaces=litgpt.model.apply_rope
)


class MergeRopeTransform(thunder.core.transform_common.Transform):
    def transform_traces_pre_prologue(self, prologue_trace, compute_trace, epilogue_trace, **kwargs):
        new_compute_trace = thunder.core.trace.from_trace(compute_trace)
        bound_symbols = compute_trace.bound_symbols[:]
        while bound_symbols:
            bsym = bound_symbols.pop(0)
            if bsym.sym == litgpt_apply_rope:
                for i, bsym2 in enumerate(bound_symbols):
                    assert not any(o is bsym.output for o in bsym2.flat_outs)
                    if bsym2.sym == litgpt_apply_rope:
                        break
                bsym2 = bound_symbols.pop(i)
                assert bsym2.sym == litgpt_apply_rope

                output = (bsym.output, bsym2.output)
                args = (bsym.args[0], bsym2.args[0], *bsym.args[1:])

                new_compute_trace.bound_symbols.append(bsym.from_bsym(args=args, output=output, sym=liger_rope))
            else:
                new_compute_trace.bound_symbols.append(bsym.from_bsym())
        new_compute_trace.set_provenance(thunder.core.trace.TraceProvenance(self.__class__))
        return prologue_trace, new_compute_trace, epilogue_trace

cfg = litgpt.Config.from_name("Llama-3.2-1B", n_layer=1)
with device:
    m = litgpt.GPT(cfg)
    m.max_seq_length = 1024
    m.set_kv_cache(1)
    inp = torch.arange(1, 6, dtype=torch.int64)[None]
    inp_pos = torch.arange(1, 6, dtype=torch.int64)


jm = thunder.jit(m, executors=(liger_ex,), transforms=(MergeRopeTransform(),))
res = jm(inp, inp_pos)

go = torch.randn_like(res)
(grad_res,) = torch.autograd.grad(res, jm.get_parameter("transformer.wte.weight"), go)
ref = m(inp, inp_pos)
(grad_ref,) = torch.autograd.grad(ref, m.get_parameter("transformer.wte.weight"), go)

assert_close(res, ref)
assert_close(grad_res, grad_ref)

assert any(bsym.sym is liger_rope_forward for bsym in thunder.last_traces(jm)[-1].bound_symbols)
assert any(bsym.sym is liger_rope_backward for bsym in thunder.last_backward_traces(jm)[-1].bound_symbols)
assert any(bsym.sym is liger_rms_norm_forward for bsym in thunder.last_traces(jm)[-1].bound_symbols)
assert any(bsym.sym is liger_rms_norm_backward for bsym in thunder.last_backward_traces(jm)[-1].bound_symbols)

if False:  # this example has additional dependencies, so we skip it in the CI
    import argparse
    import math
    import os
    from dataclasses import _MISSING_TYPE, dataclass
    import litgpt

    import datasets
    import lightning.pytorch as pl
    import torch
    import transformers
    from lightning.pytorch.strategies import DeepSpeedStrategy, FSDPStrategy
    from torch.distributed.fsdp import BackwardPrefetch, MixedPrecision
    from torch.utils.data import DataLoader
    from trl import DataCollatorForCompletionOnlyLM
    import warnings

    warnings.simplefilter(action="ignore", category=FutureWarning)


    _RETAIN_COLUMNS = {"input_ids", "attention_mask", "labels"}
    QUESTION = "<Question>"
    CHOICES = "<Choices>"


    @dataclass
    class Args:
        model: str = "meta-llama/Llama-3.2-1B-Instruct"
        data: str = "cais/mmlu"
        output_dir: str = "mmlu_finetuning"
        max_length: int = 2048
        # for llam3 8B model, deepspeed will OOM with 16 on 8XA100 80G and 8 will OOM with 8XA100 40G
        batch_size: int = 4
        lr: float = 6e-6
        weight_decay: float = 0.05
        warmup_ratio: float = 0.1
        seed: int = 42
        strategy: str = "auto"
        num_gpu: int = 1


    def warmup_cosine_schedule(warmup_steps, total_steps, min_lr=0):
        def lr_lambda(current_step):
            if current_step < warmup_steps:
                # Linear warmup
                return float(current_step) / float(max(1, warmup_steps))
            else:
                # Cosine annealing
                progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
                return max(min_lr, 0.5 * (1 + math.cos(math.pi * progress)))

        return lr_lambda


    def parse_args() -> Args:
        parser = argparse.ArgumentParser()
        for k, v in Args.__dataclass_fields__.items():
            parser.add_argument(f"--{k}", type=v.type, default=v.default)
        parsed = parser.parse_args([])
        return Args(**{k: v for k, v in vars(parsed).items() if not isinstance(v, _MISSING_TYPE)})


    class LanguageModel(pl.LightningModule):
        def __init__(self, args: Args, tokenizer):
            super().__init__()
            self.args = args
            self.tokenizer = tokenizer
            self.model = None

        def configure_model(self):
            # https://lightning.ai/docs/pytorch/stable/advanced/model_parallel/fsdp.html#speed-up-model-initialization
            if self.model is not None:
                return
            self.model = GPTForFineTuningLastToken.from_name(self.args.model.rsplit("/", 1)[-1]).to(torch.bfloat16)
            self.model.load_state_dict(litgpt.utils.lazy_load(f"checkpoints/{self.args.model}/lit_model.pth"))
            self.model = thunder.jit(
                self.model,
                executors=(liger_ex, *thunder.get_default_executors()),
                transforms=(MergeRopeTransform(), FuseSwigLUTransform(), FuseLinearCrossEntropyTransform()),
            )

        def forward(self, input_ids, attention_mask, labels=None, **kwargs):
            return self.model(idx=input_ids, mask=attention_mask, labels=labels, **kwargs)

        def training_step(self, batch):
            outputs = self.model(
                idx=batch["input_ids"],
                mask=batch["attention_mask"],
                labels=batch["labels"][:, -1],
            )
            loss = outputs
            self.log_dict(
                {"train_loss": loss},
                on_step=True,
                on_epoch=True,
                prog_bar=True,
                logger=True,
                rank_zero_only=True,
                sync_dist=False,
            )
            return loss

        def validation_step(self, batch):
            outputs = self.model(
                idx=batch["input_ids"],
                mask=batch["attention_mask"],
                labels=batch["labels"][:, -1],
            )
            loss = outputs
            self.log_dict(
                {"val_loss": loss},
                on_step=True,
                on_epoch=True,
                prog_bar=True,
                logger=True,
                rank_zero_only=True,
                sync_dist=True,
            )
            return loss

        def configure_optimizers(self):
            optimizer = torch.optim.AdamW(
                self.parameters(),
                lr=self.args.lr,
                weight_decay=self.args.weight_decay,
                fused=True,
            )
            lr_lambda = warmup_cosine_schedule(
                warmup_steps=self.trainer.estimated_stepping_batches * self.args.warmup_ratio,
                total_steps=self.trainer.estimated_stepping_batches,
                min_lr=0,
            )
            lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
            return {
                "optimizer": optimizer,
                "lr_scheduler": {"scheduler": lr_scheduler, "interval": "step"},
            }


    class DataModule(pl.LightningDataModule):
        def __init__(self, tokenizer, args: Args):
            super().__init__()
            self.train_dataset = None
            self.args = args
            self.tokenizer = tokenizer
            self.response_template_str = " <Answer>"
            response_prompt = tokenizer.encode(f"{self.response_template_str}", add_special_tokens=False)
            self.collator = DataCollatorForCompletionOnlyLM(
                tokenizer=tokenizer,
                response_template=response_prompt,
                pad_to_multiple_of=16,
            )

        def formatting_func(self, example):
            output_texts = []
            for i in range(len(example["question"])):
                choices = ""
                for j in range(len(example["choices"][i])):
                    choices += f"{j+1}. {example['choices'][i][j]}; "
                s = "Below is a question and multiple choice answers, choices separated by a semicolon. Please select the best answer for the question. "
                s += f"{QUESTION}{example['question'][i]} "
                s += f"{CHOICES}{choices} "
                s += f"{self.response_template_str}{example['answer'][i]}"
                output_texts.append(s)
            return output_texts

        def tokenize(self, example):
            outputs = self.tokenizer(
                self.formatting_func(example),
                truncation=True,
                padding=False,
                max_length=self.args.max_length,
            )
            return {
                "input_ids": outputs["input_ids"],
                "attention_mask": outputs["attention_mask"],
            }

        def setup(self, stage) -> None:
            if self.train_dataset is not None:
                return
            dataset = datasets.load_dataset(self.args.data, "auxiliary_train")
            flattened_data = [
                {
                    "answer": x["train"]["answer"],
                    "choices": x["train"]["choices"],
                    "question": x["train"]["question"],
                    "subject": x["train"]["subject"],
                }
                for x in dataset["train"]
            ][:32]
            dataset = datasets.Dataset.from_list(flattened_data)
            dataset = dataset.train_test_split(test_size=4, seed=self.args.seed)
            train_dataset, val_dataset = dataset["train"], dataset["test"]
            self.train_dataset = train_dataset.map(
                self.tokenize,
                remove_columns=list(set(train_dataset.column_names) - _RETAIN_COLUMNS),
                batched=True,
                batch_size=1,
                num_proc=4,
            )
            self.val_dataset = val_dataset.map(
                self.tokenize,
                remove_columns=list(set(val_dataset.column_names) - _RETAIN_COLUMNS),
                batched=True,
                batch_size=1,
                num_proc=4,
            )

        def train_dataloader(self):
            return DataLoader(
                self.train_dataset,
                batch_size=self.args.batch_size,
                collate_fn=self.collator,
            )

        def val_dataloader(self):
            return DataLoader(
                self.val_dataset,
                batch_size=self.args.batch_size,
                collate_fn=self.collator,
            )


    args = parse_args()
    pl.seed_everything(args.seed)
    os.makedirs(args.output_dir, exist_ok=True)

    if args.strategy == "fsdp":
        strategy = FSDPStrategy(
            auto_wrap_policy=layers,
            sharding_strategy="FULL_SHARD",
            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
            sync_module_states=True,
            activation_checkpointing_policy=layers,
            mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16),
            forward_prefetch=True,
        )
        precision = None
    elif args.strategy == "deepspeed":
        strategy = DeepSpeedStrategy(stage=3)
        precision = "bf16-mixed"
    elif args.strategy == "ddp":
        strategy = "ddp"
        precision = "bf16-true"
    else:
        strategy = "auto"
        precision = "bf16-true"

    # This only works if you have a snapshot to work from.
    trainer = pl.Trainer(
        accelerator="cuda",
        strategy=strategy,
        devices=torch.cuda.device_count() if args.num_gpu is None else args.num_gpu,
        default_root_dir=args.output_dir,
        log_every_n_steps=1,
        max_epochs=1,
        precision=precision,
    )

    tokenizer = transformers.AutoTokenizer.from_pretrained(args.model, padding_side="left", truncation_side="left")
    tokenizer.pad_token = tokenizer.eos_token
    data_module = DataModule(
        tokenizer=tokenizer,
        args=args,
    )

    model = LanguageModel(args=args, tokenizer=tokenizer)
    trainer.fit(model, datamodule=data_module)