import argparse
import collections
import json
import logging
import math
import os
import sys
import time
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F

_root = Path(__file__).parent.parent
if str(_root) not in sys.path:
    sys.path.insert(0, str(_root))

import model_cpu_gpt2 as _m

_m._fla_available = False
from model_cpu_gpt2 import CPUGPT, CPUGPTConfig

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)


def fela_250m_config() -> CPUGPTConfig:
    return CPUGPTConfig(
        n_layer=20,
        n_embd=1024,
        n_head=16,
        ffn_hidden=2752,
        fno_modes=256,
        gla_chunk=256,
        seq_len=8192,
        layer_pattern="SSSL",
        vocab_size=50257,
        bias=False,
    )


polar_express_coeffs = [
    (8.156554524902461, -22.48329292557795, 15.878769915207462),
    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
]


def adamw_step(p, grad, m, v, step, lr, b1, b2, eps, wd):
    p.mul_(1 - lr * wd)
    m.lerp_(grad, 1 - b1)
    v.lerp_(grad.square(), 1 - b2)
    bc1 = 1 - b1**step
    bc2 = 1 - b2**step
    p.addcdiv_(m / bc1, (v / bc2).sqrt_().add_(eps), value=-lr)


def muon_step(grads_stack, params, mom_buf, lr, momentum=0.95, ns_steps=3):
    mom_buf.lerp_(grads_stack, 1 - momentum)
    X = mom_buf.float()
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6)
    for a, b, c in polar_express_coeffs[:ns_steps]:
        if X.size(-2) >= X.size(-1):
            A = X.mT @ X
            X = a * X + X @ (b * A + c * (A @ A))
        else:
            A = X @ X.mT
            X = a * X + (b * A + c * (A @ A)) @ X
    torch._foreach_sub_(params, list((X * lr).to(params[0].dtype).unbind(0)))


class MuonAdamW(torch.optim.Optimizer):
    def __init__(self, param_groups):
        super().__init__(param_groups, defaults={})

    @torch.no_grad()
    def step(self):
        for g in self.param_groups:
            if g["kind"] == "adamw":
                for p in g["params"]:
                    if p.grad is None:
                        continue
                    st = self.state[p]
                    if not st:
                        st["step"] = 0
                        st["m"] = torch.zeros_like(p)
                        st["v"] = torch.zeros_like(p)
                    st["step"] += 1
                    adamw_step(
                        p,
                        p.grad,
                        st["m"],
                        st["v"],
                        st["step"],
                        g["lr"],
                        *g["betas"],
                        g["eps"],
                        g.get("wd", 0),
                    )
            elif g["kind"] == "muon":
                params = g["params"]
                if not params:
                    continue
                p0 = params[0]
                st = self.state[p0]
                stacked = torch.stack(
                    [p.grad for p in params if p.grad is not None]
                ).float()
                if not st:
                    st["mom"] = torch.zeros_like(stacked)
                lr = g["lr"] * max(1.0, p0.shape[-2] / max(p0.shape[-1], 1)) ** 0.5
                muon_step(stacked, params, st["mom"], lr, g.get("momentum", 0.95))


def build_optimizer(model, cfg, lr_matrix=0.01, lr_emb=0.02, lr_lm=0.002):
    scale = (cfg.n_embd / 768) ** -0.5
    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    matrix_params, scalar_params = [], []
    for block in raw.blocks:
        for p in block.parameters():
            (matrix_params if p.ndim == 2 else scalar_params).append(p)
    groups = [
        dict(
            kind="adamw",
            params=list(raw.wte.parameters()),
            lr=lr_emb * scale,
            betas=(0.8, 0.95),
            eps=1e-8,
        ),
        dict(
            kind="adamw",
            params=list(raw.lm_head.parameters())
            if raw.lm_head.weight is not raw.wte.weight
            else [],
            lr=lr_lm * scale,
            betas=(0.8, 0.95),
            eps=1e-8,
        ),
        dict(
            kind="adamw",
            params=scalar_params,
            lr=lr_matrix * scale,
            betas=(0.8, 0.95),
            eps=1e-8,
        ),
    ]
    for shape in sorted({p.shape for p in matrix_params}):
        ps = [p for p in matrix_params if p.shape == shape]
        groups.append(dict(kind="muon", params=ps, lr=lr_matrix, momentum=0.95))
    opt = MuonAdamW(groups)
    for g in opt.param_groups:
        g["initial_lr"] = g["lr"]
    return opt


def lr_multiplier(progress, warmup=0.01, min_ratio=0.1):
    if progress < warmup:
        return progress / warmup
    t = (progress - warmup) / (1.0 - warmup)
    return min_ratio + (1.0 - min_ratio) * 0.5 * (1.0 + math.cos(math.pi * t))


def make_loader(data_dir, seq_len, device, batch_size=1, rank=0, world_size=1):
    import glob
    import queue
    import threading

    all_parquet = sorted(glob.glob(os.path.join(data_dir, "*.parquet")))
    all_bins = sorted(glob.glob(os.path.join(data_dir, "*.bin")))

    if all_parquet:
        import pyarrow.parquet as pq

        shards = all_parquet[rank::world_size] or all_parquet

        def load_fn(path):
            tbl = pq.read_table(path, columns=["tokens"])
            t = torch.tensor(tbl["tokens"].to_pylist()[0], dtype=torch.long)
            n = (len(t) // (seq_len + 1)) * (seq_len + 1)
            return t[:n].view(-1, seq_len + 1)
    elif all_bins:
        import numpy as np

        shards = all_bins[rank::world_size] or all_bins

        def load_fn(path):
            data = np.memmap(path, dtype=np.int32, mode="r")
            n = (len(data) // (seq_len + 1)) * (seq_len + 1)
            return torch.from_numpy(data[:n].copy()).long().view(-1, seq_len + 1)
    else:
        raise FileNotFoundError(f"No parquet or bin shards in {data_dir}")

    buf: queue.Queue = queue.Queue(maxsize=4)

    def _producer():
        ep = 1
        while True:
            for shard in shards:
                seqs = load_fn(shard)
                if len(seqs) == 0:
                    continue
                idx = torch.randperm(len(seqs))
                batch = []
                for i in idx:
                    batch.append(seqs[i])
                    if len(batch) == batch_size:
                        chunk = torch.stack(batch)
                        buf.put((chunk[:, :-1], chunk[:, 1:], ep))
                        batch = []
            ep += 1

    threading.Thread(target=_producer, daemon=True).start()
    while True:
        x, y, ep = buf.get()
        yield x.to(device), y.to(device), ep


def eval_wikitext_bpb(model, cfg, device):
    import tiktoken

    try:
        from datasets import load_dataset

        try:
            ds = load_dataset(
                "Salesforce/wikitext",
                "wikitext-103-raw-v1",
                split="test",
                trust_remote_code=True,
            )
        except Exception:
            ds = load_dataset(
                "wikitext", "wikitext-103-raw-v1", split="test", trust_remote_code=True
            )
        text = "\n\n".join(ds["text"])
    except Exception as e:
        log.warning(f"WikiText load failed: {e}")
        return {"bpb": float("nan"), "nats": float("nan")}

    enc = tiktoken.get_encoding("gpt2")
    tokens = torch.tensor(enc.encode(text[:5_000_000]), dtype=torch.long)[:1_000_000]

    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    raw.eval()
    seq_len = cfg.seq_len
    n_chunks = min(100, (len(tokens) - 1) // seq_len)
    total_nll, total_tok = 0.0, 0

    with torch.no_grad():
        for i in range(n_chunks):
            x = tokens[i * seq_len : (i + 1) * seq_len].unsqueeze(0).to(device)
            y = tokens[i * seq_len + 1 : (i + 1) * seq_len + 1].unsqueeze(0).to(device)
            if x.shape[1] < seq_len:
                break
            pad = cfg.gla_chunk - (x.shape[1] % cfg.gla_chunk or cfg.gla_chunk)
            if pad < cfg.gla_chunk:
                x = F.pad(x, (0, pad))
            logits = raw(x)
            lgt = logits[0, :seq_len].float()
            loss = F.cross_entropy(lgt, y[0], reduction="sum")
            total_nll += loss.item()
            total_tok += seq_len

    nats = total_nll / max(total_tok, 1)
    bpb = nats / math.log(2)
    ppl = math.exp(nats)
    log.info(f"WikiText eval: BPB={bpb:.4f}  nats={nats:.4f}  PPL={ppl:.2f}")
    return {"bpb": bpb, "nats": nats, "ppl": ppl, "tokens_evaluated": total_tok}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", default="/data/openwebtext")
    parser.add_argument("--out-dir", default="results/scaling_midpoint_v2")
    parser.add_argument("--device-batch", type=int, default=4)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--resume", type=str, default=None)
    args = parser.parse_args()

    use_ddp = "RANK" in os.environ
    local_rank, global_rank = 0, 0
    world_size = 1
    if use_ddp:
        dist.init_process_group(backend="nccl")
        local_rank = dist.get_rank() % torch.cuda.device_count()
        global_rank = dist.get_rank()
        world_size = dist.get_world_size()
        torch.cuda.set_device(local_rank)

    device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
    is_master = global_rank == 0
    num_gpus = world_size

    cfg = fela_250m_config()
    model = CPUGPT(cfg).to(device)
    if use_ddp:
        model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

    raw = model.module if use_ddp else model
    actual_params = raw.param_count()

    if is_master:
        log.info(f"FELA ~250M v2: {actual_params / 1e6:.1f}M params")
        log.info(f"  n_layer={cfg.n_layer}  n_embd={cfg.n_embd}  seq_len={cfg.seq_len}")
        log.info(f"  lr_emb=0.02  lr_matrix=0.01  warmup=1%  [FIXED from v1]")

    opt = build_optimizer(model, cfg)

    chinchilla_tokens = actual_params * 20
    if args.smoke:
        chinchilla_tokens = 10_000_000
        log.info("Smoke: 10M tokens")
    else:
        if is_master:
            log.info(f"Chinchilla target: {chinchilla_tokens / 1e9:.2f}B tokens")

    total_batch_tokens = args.device_batch * cfg.seq_len * num_gpus
    total_steps = max(1, int(chinchilla_tokens / total_batch_tokens))

    if is_master:
        log.info(f"  total_steps={total_steps}  batch_tokens={total_batch_tokens:,}")

    out_dir = Path(args.out_dir)
    ckpt_dir = out_dir / "ckpt"
    if is_master:
        ckpt_dir.mkdir(parents=True, exist_ok=True)

    start_step = 0
    if args.resume:
        ck = torch.load(args.resume, map_location="cpu")
        sd = {
            k.replace("module.", "").replace("_orig_mod.", ""): v
            for k, v in ck["model"].items()
        }
        raw.load_state_dict(sd)
        try:
            opt.load_state_dict(ck["opt"])
        except Exception:
            pass
        start_step = ck.get("step", 0)
        if is_master:
            log.info(f"Resumed from {args.resume} at step {start_step}")

    loader = make_loader(
        args.data_dir, cfg.seq_len, device, args.device_batch, global_rank, world_size
    )

    loss_window: collections.deque = collections.deque(maxlen=50)
    t0 = time.perf_counter()
    step = start_step

    for x, y, _ in loader:
        if step >= total_steps:
            break

        model.train()
        progress = step / total_steps
        mult = lr_multiplier(progress)
        for g in opt.param_groups:
            g["lr"] = g["initial_lr"] * mult

        with torch.autocast(device_type=device.split(":")[0], dtype=torch.bfloat16):
            loss = model(x, y)

        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        opt.zero_grad(set_to_none=True)
        step += 1

        raw_loss = loss.item()
        loss_window.append(raw_loss)

        if is_master and step % 100 == 0:
            elapsed = time.perf_counter() - t0
            tps = (step - start_step) * total_batch_tokens / elapsed
            eta_h = (total_steps - step) * total_batch_tokens / max(tps, 1) / 3600
            avg_loss = sum(loss_window) / len(loss_window)
            log.info(
                f"step={step}/{total_steps} ({100 * step / total_steps:.0f}%)  "
                f"loss={raw_loss:.4f}  avg50={avg_loss:.4f}  "
                f"tok/s={tps:,.0f}  ETA={eta_h:.1f}h"
            )

            if step >= 10000 and avg_loss > 6.0:
                log.error(
                    f"DIVERGENCE DETECTED: avg50={avg_loss:.4f} > 6.0 at step {step}"
                )
                log.error("Aborting. Check LR config or data.")
                if use_ddp:
                    dist.destroy_process_group()
                sys.exit(1)

        if is_master and step % 5000 == 0:
            ck_path = ckpt_dir / f"step_{step:06d}.pt"
            torch.save(
                {"step": step, "model": raw.state_dict(), "opt": opt.state_dict()},
                ck_path,
            )
            avg_loss = sum(loss_window) / len(loss_window)
            log.info(f"Checkpoint saved: {ck_path}  avg50_loss={avg_loss:.4f}")

    if is_master:
        torch.save(
            {"step": step, "model": raw.state_dict(), "opt": opt.state_dict()},
            ckpt_dir / "final.pt",
        )
        log.info(f"Training complete at step {step}. Saved final.pt")

    if use_ddp:
        dist.barrier()

    if is_master:
        log.info("=== WikiText BPB eval ===")
        raw.eval()
        wt = eval_wikitext_bpb(raw, cfg, device)

        tokens_trained = step * total_batch_tokens
        summary = {
            "model": "FELA_250M_v2",
            "params_M": actual_params / 1e6,
            "n_layer": cfg.n_layer,
            "n_embd": cfg.n_embd,
            "seq_len": cfg.seq_len,
            "tokens_trained": tokens_trained,
            "chinchilla_optimal_tokens": chinchilla_tokens,
            "pct_chinchilla": round(tokens_trained / chinchilla_tokens * 100, 1),
            "steps": step,
            "wikitext_bpb": wt["bpb"],
            "wikitext_nats": wt.get("nats"),
            "wikitext_ppl": wt.get("ppl"),
            "lr_emb": 0.02,
            "lr_matrix": 0.01,
            "fix_note": "v1 had lr_emb=0.1 (5x too high), caused divergence at step 25K",
        }

        (out_dir / "wikitext.json").write_text(json.dumps(wt, indent=2))
        (out_dir / "summary.json").write_text(json.dumps(summary, indent=2))

        print("\n" + "=" * 60)
        print("SCALING LAW DATA POINTS")
        print("=" * 60)
        print(f"  {'Model':<20} {'Params':>8}  {'Tokens':>8}  {'BPB':>8}")
        print(f"  {'-' * 48}")
        print(f"  {'FELA 124M (v9)':<20} {'124M':>8}  {'2B':>8}  {'1.297':>8}")
        print(
            f"  {'FELA ~250M (v2)':<20} {actual_params / 1e6:>7.0f}M  "
            f"{tokens_trained / 1e9:>7.1f}B  {wt['bpb']:>8.4f}"
        )
        print(f"  {'FELA 1.13B (22B)':<20} {'1130M':>8}  {'22B':>8}  {'TBD':>8}")
        print("=" * 60)
        log.info(f"Outputs: {out_dir}/")

    if use_ddp:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()