#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from functools import partial
from pathlib import Path

import torch
import torch.nn.functional as F

from megatron.bridge.models.gpt_provider import GPTModelProvider
from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
from megatron.bridge.recipes.common import _pretrain_common
from megatron.bridge.training.config import ConfigContainer
from megatron.bridge.training.gpt_step import forward_step
from megatron.bridge.training.pretrain import pretrain


REPO_ROOT = Path(__file__).resolve().parents[2]
TOKENIZER_DIR = REPO_ROOT / "tokenizer/glm5.2"
DEFAULT_VALIDATION_PREFIX = REPO_ROOT / "dataset/val/megatron_8192_glm52/heldout_2p8k_text_document"


TARGET_CATEGORY_MIX = {
    "english_web": 0.40,
    "english_edu": 0.20,
    "chinese_clean": 0.10,
    "science": 0.10,
    "logic": 0.10,
    "math": 0.05,
    "code": 0.05,
}


def _normalize_weights(raw_weights: list[float]) -> list[float]:
    total = sum(raw_weights)
    if total <= 0:
        raise ValueError("Data blend weights sum to zero.")
    return [weight / total for weight in raw_weights]


def _token_weights_from_manifest(prefix_list: list[str], manifest: dict) -> list[float]:
    token_by_prefix = {}
    for result in manifest.get("results") or []:
        output_prefix = result.get("output_prefix")
        tokens = result.get("tokens_estimated_from_int32_bin")
        if output_prefix and tokens:
            token_by_prefix[str(output_prefix) + "_text_document"] = float(tokens)

    weights = [token_by_prefix.get(prefix, 0.0) for prefix in prefix_list]
    if any(weight <= 0 for weight in weights):
        missing = [prefix for prefix, weight in zip(prefix_list, weights) if weight <= 0]
        raise ValueError(f"Missing token counts for {len(missing)} data prefixes; first={missing[:3]}")
    return _normalize_weights(weights)


def _target_mix_weights_from_category_stats(
    prefix_list: list[str],
    category_stats_path: Path,
) -> list[float] | None:
    if not category_stats_path.exists():
        return None

    stats = json.loads(category_stats_path.read_text(encoding="utf-8"))
    per_prefix = stats.get("prefix_category_tokens") or {}
    category_totals = stats.get("category_tokens") or {}
    raw_weights = []
    for prefix in prefix_list:
        category_tokens = per_prefix.get(prefix)
        if not category_tokens:
            return None

        shard_weight = 0.0
        for category, tokens in category_tokens.items():
            target_ratio = TARGET_CATEGORY_MIX.get(category, 0.0)
            category_total = float(category_totals.get(category, 0.0))
            if target_ratio > 0 and category_total > 0 and tokens:
                shard_weight += target_ratio * (float(tokens) / category_total)
        raw_weights.append(shard_weight)

    return _normalize_weights(raw_weights)


def load_data_blend(args: argparse.Namespace) -> tuple[list[str], list[float] | None] | None:
    if args.data_manifest:
        manifest_path = Path(args.data_manifest)
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        prefixes = manifest.get("ok_prefixes") or []
        if not prefixes:
            raise ValueError(f"{manifest_path} has no ok_prefixes")
        prefix_list = [str(prefix) for prefix in prefixes]
        category_stats_path = manifest_path.with_name("prefix_category_stats.json")
        weights = _target_mix_weights_from_category_stats(prefix_list, category_stats_path)
        if weights is None:
            weights = _token_weights_from_manifest(prefix_list, manifest)
        return (prefix_list, weights)
    if args.data_prefix:
        return ([args.data_prefix], None)
    return None


def build_config(args: argparse.Namespace) -> ConfigContainer:
    cfg = _pretrain_common()

    cfg.model = GPTModelProvider(
        num_layers=12,
        hidden_size=1536,
        num_attention_heads=24,
        num_query_groups=4,
        ffn_hidden_size=4608,
        seq_length=args.seq_length,
        vocab_size=154856,
        should_pad_vocab=True,
        share_embeddings_and_output_weights=True,
        position_embedding_type="rope",
        normalization="RMSNorm",
        gated_linear_unit=True,
        activation_func=F.silu,
        add_bias_linear=False,
        num_moe_experts=12,
        moe_layer_freq=[1 if idx in {2, 4, 6, 8, 10} else 0 for idx in range(12)],
        moe_ffn_hidden_size=6144,
        moe_router_topk=4,
        moe_router_load_balancing_type="aux_loss",
        moe_aux_loss_coeff=0.02,
        moe_z_loss_coeff=0.001,
        moe_token_dispatcher_type="alltoall",
        moe_expert_capacity_factor=1.25,
        moe_router_enable_expert_bias=False,
        moe_router_bias_update_rate=0.02,
        moe_grouped_gemm=not args.disable_moe_grouped_gemm,
        tensor_model_parallel_size=args.tensor_parallel,
        pipeline_model_parallel_size=args.pipeline_parallel,
        expert_model_parallel_size=args.expert_parallel,
        context_parallel_size=args.context_parallel,
        sequence_parallel=args.tensor_parallel > 1,
        transformer_impl="transformer_engine",
        attention_backend="flash",
        transformer_layer_spec=partial(get_gpt_decoder_block_spec, use_transformer_engine=True),
        init_method_std=0.02,
        recompute_granularity=args.recompute_granularity,
        recompute_method=args.recompute_method,
        recompute_num_layers=args.recompute_num_layers,
        recompute_modules=args.recompute_modules,
    )

    cfg.tokenizer.tokenizer_type = "HuggingFaceTokenizer"
    cfg.tokenizer.tokenizer_model = str(TOKENIZER_DIR)
    cfg.dataset.seq_length = args.seq_length
    cfg.dataset.split = args.split
    cfg.dataset.mid_level_dataset_surplus = 0.05
    data_blend = load_data_blend(args)
    if data_blend:
        validation_prefix = Path(args.validation_prefix) if args.validation_prefix else None
        if validation_prefix and validation_prefix.with_suffix(".bin").exists() and validation_prefix.with_suffix(".idx").exists():
            cfg.dataset.data_path = None
            cfg.dataset.split = None
            cfg.dataset.blend = None
            cfg.dataset.blend_per_split = [
                data_blend,
                ([str(validation_prefix)], None),
                None,
            ]
        else:
            cfg.dataset.data_path = data_blend[0][0]
            cfg.dataset.blend = data_blend
    else:
        cfg.dataset.blend = None

    cfg.train.train_iters = args.train_iters
    cfg.train.micro_batch_size = args.micro_batch_size
    cfg.train.global_batch_size = args.global_batch_size
    cfg.optimizer.lr = args.lr
    cfg.optimizer.min_lr = args.min_lr
    cfg.optimizer.weight_decay = args.weight_decay
    cfg.scheduler.lr_warmup_fraction = args.warmup_fraction
    cfg.scheduler.lr_warmup_iters = 0
    cfg.scheduler.lr_warmup_samples = 0

    cfg.checkpoint.save = args.save_dir
    cfg.checkpoint.load = args.load_dir
    cfg.checkpoint.save_interval = args.save_interval
    cfg.checkpoint.most_recent_k = args.keep_last_checkpoints

    cfg.logger.tensorboard_dir = args.tensorboard_dir
    cfg.logger.log_interval = args.log_interval
    cfg.logger.log_progress = True
    cfg.logger.log_runtime_to_tensorboard = True
    cfg.logger.log_throughput = True
    cfg.logger.log_throughput_to_tensorboard = True
    cfg.logger.log_validation_ppl_to_tensorboard = True
    cfg.validation.eval_interval = args.eval_interval
    cfg.validation.eval_iters = args.eval_iters

    cfg.ddp.use_megatron_fsdp = False
    cfg.ddp.overlap_grad_reduce = args.overlap_grad_reduce
    cfg.ddp.overlap_param_gather = args.overlap_param_gather
    cfg.ddp.check_for_nan_in_grad = True
    cfg.ddp.use_distributed_optimizer = args.use_distributed_optimizer
    cfg.optimizer.use_distributed_optimizer = args.use_distributed_optimizer
    cfg.optimizer.overlap_param_gather = args.overlap_param_gather

    return cfg


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Laoyao 2B MoE Megatron-Bridge pretrain launcher.")
    parser.add_argument("--data-prefix", default=None, help="Megatron indexed dataset prefix without .bin/.idx.")
    parser.add_argument("--data-manifest", default=None, help="Manifest produced by convert_pretrain_parquet_to_megatron.py.")
    parser.add_argument(
        "--validation-prefix",
        default=str(DEFAULT_VALIDATION_PREFIX),
        help="Optional Megatron indexed dataset prefix used as the validation split.",
    )
    parser.add_argument("--seq-length", type=int, default=8192)
    parser.add_argument("--train-iters", type=int, default=10)
    parser.add_argument("--micro-batch-size", type=int, default=1)
    parser.add_argument("--global-batch-size", type=int, default=1024)
    parser.add_argument("--tensor-parallel", type=int, default=1)
    parser.add_argument("--pipeline-parallel", type=int, default=1)
    parser.add_argument("--expert-parallel", type=int, default=1)
    parser.add_argument("--context-parallel", type=int, default=1)
    parser.add_argument("--disable-moe-grouped-gemm", action="store_true", help="Disable grouped GEMM for MoE experts.")
    parser.add_argument("--recompute-granularity", choices=["full", "selective"], default=None)
    parser.add_argument("--recompute-method", choices=["uniform", "block"], default=None)
    parser.add_argument("--recompute-num-layers", type=int, default=None)
    parser.add_argument("--recompute-modules", nargs="*", default=None)
    parser.add_argument("--use-distributed-optimizer", action="store_true", help="Enable Megatron distributed optimizer.")
    parser.add_argument("--overlap-grad-reduce", action="store_true", help="Overlap DDP gradient reduction.")
    parser.add_argument("--overlap-param-gather", action="store_true", help="Overlap distributed optimizer parameter gather.")
    parser.add_argument("--dataset-workers", type=int, default=8)
    parser.add_argument("--split", default="9999,8,2")
    parser.add_argument("--lr", type=float, default=3.0e-4)
    parser.add_argument("--min-lr", type=float, default=5.0e-6)
    parser.add_argument("--weight-decay", type=float, default=0.1)
    parser.add_argument("--warmup-fraction", type=float, default=0.01)
    parser.add_argument("--save-dir", default="/mnt/beegfs/yi/laoyao_2b_moe/runs/megatron_bridge/checkpoints")
    parser.add_argument("--load-dir", default=None)
    parser.add_argument("--tensorboard-dir", default="/mnt/beegfs/yi/laoyao_2b_moe/runs/megatron_bridge/tensorboard")
    parser.add_argument("--save-interval", type=int, default=1000)
    parser.add_argument("--keep-last-checkpoints", type=int, default=10)
    parser.add_argument("--log-interval", type=int, default=10)
    parser.add_argument("--eval-interval", type=int, default=1000)
    parser.add_argument("--eval-iters", type=int, default=10)
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    cfg = build_config(args)
    if args.dry_run:
        print("laoyao_megatron_bridge_config_ok")
        print(f"seq_length={cfg.model.seq_length}")
        print(f"hidden_size={cfg.model.hidden_size}")
        print(f"num_moe_experts={cfg.model.num_moe_experts}")
        print(f"moe_router_topk={cfg.model.moe_router_topk}")
        print(f"moe_layer_freq={cfg.model.moe_layer_freq}")
        print(f"moe_router_load_balancing_type={cfg.model.moe_router_load_balancing_type}")
        print(f"moe_aux_loss_coeff={cfg.model.moe_aux_loss_coeff}")
        print(f"moe_grouped_gemm={cfg.model.moe_grouped_gemm}")
        print(f"recompute_granularity={cfg.model.recompute_granularity}")
        print(f"recompute_method={cfg.model.recompute_method}")
        print(f"recompute_num_layers={cfg.model.recompute_num_layers}")
        print(f"recompute_modules={cfg.model.recompute_modules}")
        print(f"use_distributed_optimizer={cfg.optimizer.use_distributed_optimizer}")
        print(f"overlap_grad_reduce={cfg.ddp.overlap_grad_reduce}")
        print(f"overlap_param_gather={cfg.ddp.overlap_param_gather}")
        print(f"tokenizer_model={cfg.tokenizer.tokenizer_model}")
        print(f"dataset_prefix={args.data_prefix or 'mock'}")
        print(f"dataset_manifest={args.data_manifest or 'none'}")
        print(f"dataset_blend_size={len(cfg.dataset.blend or [])}")
        print(f"dataset_blend_per_split={cfg.dataset.blend_per_split is not None}")
        return
    pretrain(config=cfg, forward_step_func=forward_step)


if __name__ == "__main__":
    main()
