import json
import re

from collections import Counter


def should_ignore(weight_name, exact_ignore: list, regex_ignore: list):
    if any(weight_name.startswith(ig) for ig in exact_ignore):
        return True
    if any(p.match(weight_name) for p in regex_ignore):
        return True
    
    return False


def compile_ignore(ignore_list: list | None = None, 
                   extra_ignore_list: list | None = None):
    exact = set()
    regex = []

    for x in (ignore_list or []) + (extra_ignore_list or []):
        if x.startswith("re:"):
            regex.append(re.compile(x[3:]))
        else:
            exact.add(x)
    return exact, regex


def normalize_to_regex(weight_name: str) -> str:
    parts = []
    for p in weight_name.split("."):
        if p.isdigit():
            parts.append(r"\d+")
        else:
            parts.append(p)
    return "re:" + ".".join(parts) + ".*"


def infer_default_quant_target(quant_config: dict) -> str:
    classes = quant_config.get("class") or {}
    int4_list = classes.get("int4") or []
    int8_list = classes.get("int8") or []

    if int4_list and not int8_list:
        return "int4"
    if int8_list and not int4_list:
        return "int8"
    if len(int4_list) >= len(int8_list):
        return "int4"
    return "int8"


def suggest_category(weight_name: str, default_quant_target: str) -> tuple[str, str]:
    wn = weight_name.lower()

    ignore_keywords = [
        "lm_head",
        "embed_tokens",
        ".input_layernorm.",
        ".post_attention_layernorm.",
        ".self_attn.k_norm.",
        ".self_attn.q_norm.",
        ".norm.",
        ".bias",
        ".e_score_correction_bias",
    ]
    if any(k in wn for k in ignore_keywords):
        return "ignore", "norm/embed/bias 类型，通常放在 ignore"

    quant_keywords = [
        ".self_attn.",
        ".mlp.",
        ".experts.",
        "_proj.",
        ".proj.",
        ".linear.",
    ]
    if any(k in wn for k in quant_keywords):
        return default_quant_target, f"线性/投影权重，建议放在 {default_quant_target}"

    return "manual-review", "命名不典型，建议人工确认后归类"


def collect_quant_ignore_list(quant_config: dict) -> list[str]:
    classes = quant_config.get("class") or {}
    ignore_list = []
    ignore_list += quant_config.get("ignore") or []
    ignore_list += classes.get("int4") or []
    ignore_list += classes.get("int8") or []
    return ignore_list


def build_missing_summary(missing: list[str], quant_config: dict, max_show: int) -> dict:
    default_quant_target = infer_default_quant_target(quant_config)
    category_counter = Counter()
    rows = []

    for wn in missing:
        category, reason = suggest_category(wn, default_quant_target)
        category_counter[category] += 1
        rows.append(
            {
                "weight_name": wn,
                "suggested_category": category,
                "suggested_pattern": normalize_to_regex(wn),
                "reason": reason,
            }
        )

    return {
        "missing_total": len(missing),
        "default_quant_target": default_quant_target,
        "category_summary": {
            "ignore": category_counter.get("ignore", 0),
            "int4": category_counter.get("int4", 0),
            "int8": category_counter.get("int8", 0),
            "manual-review": category_counter.get("manual-review", 0),
        },
        "top_missing": rows[:max_show],
    }


def summarize_missing(missing: list[str], quant_config: dict, max_show: int):
    summary = build_missing_summary(missing=missing, quant_config=quant_config, max_show=max_show)
    print(f"Missing total: {len(missing)}")
    print(f"Default quant target inferred from config: {summary['default_quant_target']}")
    print("Suggested category summary:")
    for cat, count in summary["category_summary"].items():
        print(f"  - {cat}: {count}")

    if not missing:
        print("All weights are covered by quant.ignore/int4/int8.")
        return

    print(f"\nTop {len(summary['top_missing'])} missing items:")
    for idx, row in enumerate(summary["top_missing"], start=1):
        print(f"[{idx}] {row['weight_name']}")
        print(f"    suggestion: {row['suggested_category']}")
        print(f"    reason: {row['reason']}")
        print(f"    pattern: {row['suggested_pattern']}")


def check(weight_names, ignore_list):
    exact_ignore, regex_ignore = compile_ignore(ignore_list)
    missing = []
    for wn in weight_names:
        if not should_ignore(wn, exact_ignore, regex_ignore):
            missing.append(wn)
    return missing


def build_validation_result(
    weight_names: list[str], quant_config: dict, max_show: int = 20
) -> dict:
    ignore_list = collect_quant_ignore_list(quant_config)
    missing = check(weight_names, ignore_list)
    return build_missing_summary(missing=missing, quant_config=quant_config, max_show=max_show)


if __name__ == "__main__":
    import yaml

    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument("--weight-map", type=str)
    parser.add_argument("--model-info", type=str)
    parser.add_argument("--max-show", type=int, default=20)

    args = parser.parse_args()

    with open(args.model_info, "r") as f:
        cfg = yaml.safe_load(f)
    
    quant_config = cfg['quant']
    ignore_list = collect_quant_ignore_list(quant_config)

    with open(args.weight_map, "r") as f:
        weight_names = json.load(f)['weight_map'].keys()

    missing = check(list(weight_names), ignore_list)
    summarize_missing(missing, quant_config, args.max_show)