# -*- coding: utf-8 -*-
"""
tools/ptv3_eval_0920.py  —— 0920 抓日志 + 统计 + PPT 卡片 + 对比卡片（base vs quant）
- 只读日志与 ckpt（不再构建模型或跑 ptflops）
- 抓取：mIoU/Acc、(可选) FLOPs/训练速度；从 ckpt 估 avg weight bit & 量化比例
- 产物：
  exp/summary_0920/summary_all_0920.csv
  exp/summary_0920/<dataset>_{base,quant}_0920.png
  exp/summary_0920/<dataset>_compare_0920.png
"""

import os, re, csv, math, argparse
from pathlib import Path
import torch

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

DATASET_KIND = {
    "scannet": "seg",
    "s3dis":   "seg",
    "nuscenes":"seg",
    "modelnet":"cls",
}

# ---------- 基础工具 ----------
def guess_dataset(run_dir: Path):
    nm = run_dir.name.lower()
    for k in DATASET_KIND:
        if k in nm:
            return k
    for p in run_dir.glob("*.log"):
        t = p.read_text(errors="ignore").lower()
        for k in DATASET_KIND:
            if k in t:
                return k
    return "unknown"

def find_logs(run_dir: Path):
    cand = []
    for n in ["train.log", "eval.log"]:
        p = run_dir / n
        if p.exists(): cand.append(p)
    cand += [p for p in run_dir.glob("*.log") if p not in cand]
    return cand

def find_ckpt(run_dir: Path):
    pats = ["model_best*.pth","best*.pth","model_last*.pth","last*.pth",
            "latest*.pth","*checkpoint*.pth","*ckpt*.pth","*.pth","*.pt"]
    hits=[]
    for p in pats:
        hits += list(run_dir.rglob(p))
    hits = [h for h in hits if h.is_file()]
    if not hits: return None
    scored=[]
    for p in hits:
        s=0; n=p.name.lower()
        if "best" in n: s+=100
        if "last" in n or "latest" in n: s+=50
        s += int(p.stat().st_mtime)
        scored.append((s,p))
    scored.sort(reverse=True)
    return scored[0][1]

def load_state_dict_any(p: Path):
    try:
        try:
            obj = torch.load(p, map_location="cpu", weights_only=True)
        except TypeError:
            obj = torch.load(p, map_location="cpu")
        if isinstance(obj, dict):
            for k in ["state_dict","model","net","module","ema","model_state","model_ema"]:
                if k in obj and isinstance(obj[k], dict):
                    return {kk:vv for kk,vv in obj[k].items() if torch.is_tensor(vv)}
            if all(isinstance(k,str) for k in obj.keys()):
                return {k:v for k,v in obj.items() if torch.is_tensor(v)}
        return {}
    except Exception:
        return {}

def avg_weight_bits(sd, w_bits=2,
                    exclude=("cls_head","embedding.stem","stem","head"),
                    excl_norm_bias=True, force_fp32=False):
    total=qcnt=fpcnt=0
    for name, t in sd.items():
        if not torch.is_tensor(t): continue
        n = t.numel(); lname = name.lower()
        excl = any(h in lname for h in exclude)
        if excl_norm_bias and (".norm" in lname or "bn" in lname or lname.endswith(".bias")):
            excl = True
        total += n
        if excl or force_fp32: fpcnt += n
        else: qcnt += n
    if total==0: return float("nan"), None
    avg = (qcnt*(32 if force_fp32 else w_bits) + fpcnt*32.0)/total
    qratio = 100.0 * (qcnt/total)
    return avg, qratio

def parse_bits_and_mode_by_name_or_log(run_dir: Path):
    name = run_dir.name.lower()
    w=a=None
    m=re.search(r"w(\d+)a(\d+)", name)
    if m: w,a=int(m.group(1)), int(m.group(2))
    mode="quant" if m else ("fp32" if ("fp32" in name or "baseline" in name) else "unknown")
    t = ""
    for lp in find_logs(run_dir):
        t += "\n" + lp.read_text(errors="ignore").lower()
    m1=re.search(r"quant\d*\.?enable\s*=\s*(true|false)", t)
    if m1: mode="quant" if m1.group(1)=="true" else "fp32"
    mw=re.search(r"quant\d*\.?w_bits\s*=\s*(\d+)", t)
    ma=re.search(r"quant\d*\.?a_bits\s*=\s*(\d+)", t)
    if mw: w=int(mw.group(1))
    if ma: a=int(ma.group(1))
    if mode=="fp32": w=w or 32; a=a or 32
    return w or 2, a or 8, mode

# ---------- 日志解析 ----------
RE_MIOU_TRIPLE = re.compile(
    r'Val\s+result\s*:\s*mIoU\s*/\s*mAcc\s*/\s*allAcc\s*([0-9]*\.?[0-9]+)\s*/\s*([0-9]*\.?[0-9]+)\s*/\s*([0-9]*\.?[0-9]+)',
    re.I
)
RE_MIOU_BEST1 = re.compile(r'Currently\s+Best\s+mIoU\s*:\s*([0-9]*\.?[0-9]+)', re.I)
RE_MIOU_BEST2 = re.compile(r'Best\s+validation\s+mIoU\s+updated\s+to\s*:\s*([0-9]*\.?[0-9]+)', re.I)
RE_MIOU_KV    = re.compile(r'(?<!/)\bmIoU\b\s*[:=]\s*([0-9]*\.?[0-9]+)', re.I)

def _pct(v):
    v=float(v)
    return v*100.0 if v<=1.2 else v

def parse_miou_from_logs(log_paths):
    vals=[]
    for p in log_paths:
        s=p.read_text(errors="ignore")
        vals += [ _pct(m.group(1)) for m in RE_MIOU_TRIPLE.finditer(s) ]
        vals += [ _pct(m.group(1)) for m in RE_MIOU_BEST1.finditer(s) ]
        vals += [ _pct(m.group(1)) for m in RE_MIOU_BEST2.finditer(s) ]
        vals += [ _pct(m.group(1)) for m in RE_MIOU_KV.finditer(s) ]
    return max(vals) if vals else None

RE_ACC_KEYS = [
    re.compile(r'\bAcc(?:_best)?\b\s*[:=]\s*([0-9]*\.?[0-9]+)%?', re.I),
    re.compile(r'\baccuracy\b\s*[:=]\s*([0-9]*\.?[0-9]+)%?', re.I),
    re.compile(r'\bOA\b\s*[:=]\s*([0-9]*\.?[0-9]+)%?', re.I),
    re.compile(r'\bTop-?1\b\s*[:=]\s*([0-9]*\.?[0-9]+)%?', re.I),
]

def parse_acc_from_logs(log_paths):
    vals=[]
    for p in log_paths:
        s=p.read_text(errors="ignore")
        for rex in RE_ACC_KEYS:
            vals += [ float(m.group(1)) if float(m.group(1))>1.2 else float(m.group(1))*100.0
                      for m in rex.finditer(s) ]
    return max(vals) if vals else None

RE_FLOPS = re.compile(
    r'(?P<val>\d+(?:\.\d+)?(?:e[+-]?\d+)?)\s*(?P<unit>GFLOPs|GFlops|GFLOP|GMACs|GMac|MACs|MAC)\b',
    re.I
)
def _to_float(x):
    try: return float(x)
    except Exception:
        try: return float(eval(x))
        except Exception: return None

def parse_flops_from_logs(log_paths):
    last=None
    for p in log_paths:
        s=p.read_text(errors="ignore")
        for m in RE_FLOPS.finditer(s):
            last=m
    if not last: return None
    val=_to_float(last.group('val')); unit=last.group('unit').lower()
    if val is None: return None
    if 'gflop' in unit: return round(val,3)
    if 'gmac'  in unit: return round(val*2.0,3)
    if 'mac'   in unit: return round((val/1e9)*2.0,3)
    return None

RE_MS_PER_IT  = re.compile(r'(?P<val>\d+(?:\.\d+)?)\s*ms\s*/\s*it', re.I)
RE_SEC_PER_IT = re.compile(r'(?P<val>\d+(?:\.\d+)?)\s*s(?:ec)?\s*/\s*it', re.I)
RE_ITS        = re.compile(r'(?P<val>\d+(?:\.\d+)?)\s*it\s*/\s*s', re.I)
RE_FPS        = re.compile(r'(?P<val>\d+(?:\.\d+)?)\s*fps\b', re.I)
RE_SPS        = re.compile(r'(?P<val>\d+(?:\.\d+)?)\s*samples\s*/\s*s', re.I)

def parse_speed_from_logs(log_paths):
    ms=None; fps=None
    for p in log_paths:
        s=p.read_text(errors="ignore")
        hits=list(RE_MS_PER_IT.finditer(s))
        if hits: ms=float(hits[-1].group('val'))
        else:
            sec=list(RE_SEC_PER_IT.finditer(s))
            if sec: ms=float(sec[-1].group('val'))*1000.0
            else:
                its=list(RE_ITS.finditer(s))
                if its:
                    v=float(its[-1].group('val'))
                    ms=1000.0/v if v>0 else None
        f=list(RE_FPS.finditer(s))
        if f: fps=float(f[-1].group('val'))
        else:
            sp=list(RE_SPS.finditer(s))
            if sp: fps=float(sp[-1].group('val'))
    return ms, fps

# ---------- 绘图 ----------
def draw_card_png(dataset, tag, task_kind, info, out_png):
    plt.figure(figsize=(10.5, 6.5))
    ax = plt.gca(); ax.axis("off")
    main_name = "mIoU (%)" if task_kind == "seg" else "Acc (%)"
    main_val  = info.get("miou") if task_kind=="seg" else info.get("acc")
    rows = [
        ["Dataset", dataset.upper()],
        ["Variant", tag.upper()],
        [main_name, "-" if main_val is None else f"{main_val:.2f}"],
        ["Avg Weight Bit", "-" if info.get("avg_bit") is None else f"{info['avg_bit']:.3f}"],
        ["Quant Ratio (%)", "-" if info.get("qratio") is None else f"{info['qratio']:.2f}"],
    ]
    if info.get("flops") is not None:
        rows.append(["FLOPs (GFLOPs)", f"{info['flops']:.2f}"])
    ms_it, fps = info.get("ms_per_it"), info.get("fps")
    if (ms_it is not None) or (fps is not None):
        desc=[]
        if ms_it is not None: desc.append(f"{ms_it:.2f} ms/it")
        if fps   is not None: desc.append(f"{fps:.2f} fps")
        rows.append(["Train speed (log)", " | ".join(desc)])
    table = ax.table(cellText=rows, colLabels=["Metric","Value"], loc="center", cellLoc="center")
    table.auto_set_font_size(False); table.set_fontsize(18); table.scale(1.4, 2.0)
    for (r,c),cell in table.get_celld().items():
        if r==0: cell.set_text_props(weight='bold')
    ax.set_title(f"{dataset.upper()} - {tag.upper()} (0920)", fontsize=28, fontweight="bold", pad=20)
    plt.tight_layout()
    Path(out_png).parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(out_png, dpi=220, bbox_inches="tight")
    plt.close()
    print(f"[card] {out_png}")

def draw_compare_png(dataset, task_kind, base_info, quant_info, out_png):
    plt.figure(figsize=(12, 6.5))
    ax = plt.gca(); ax.axis("off")
    metric = "mIoU (%)" if task_kind=="seg" else "Acc (%)"
    b = base_info.get("miou") if task_kind=="seg" else base_info.get("acc")
    q = quant_info.get("miou") if task_kind=="seg" else quant_info.get("acc")
    delta = None if (b is None or q is None) else (q - b)

    def fmt_bit(v): return "-" if v is None else f"{v:.3f}"
    bit_base  = fmt_bit(base_info.get("avg_bit"))
    bit_quant = fmt_bit(quant_info.get("avg_bit"))
    qratio_s  = "-" if quant_info.get("qratio") is None else f"{quant_info['qratio']:.2f}"

    rows = [
        ["Dataset", dataset.upper()],
        [metric + " (BASE)", "-" if b is None else f"{b:.2f}"],
        [metric + " (QUANT)", "-" if q is None else f"{q:.2f}"],
        [f"Δ {metric}", "-" if delta is None else f"{delta:+.2f}"],
        ["Avg Bit (BASE/QUANT)", f"{bit_base} / {bit_quant}"],
        ["Quant Ratio (QUANT %)", qratio_s],
    ]
    table = ax.table(cellText=rows, colLabels=["Metric","Value"], loc="center", cellLoc="center")
    table.auto_set_font_size(False); table.set_fontsize(18); table.scale(1.6, 2.0)
    for (r,c),cell in table.get_celld().items():
        if r==0: cell.set_text_props(weight='bold')
    ax.set_title(f"{dataset.upper()} - BASE vs QUANT (0920)", fontsize=28, fontweight="bold", pad=20)
    plt.tight_layout()
    Path(out_png).parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(out_png, dpi=220, bbox_inches="tight")
    plt.close()
    print(f"[compare] {out_png}")

# ---------- 主流程 ----------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--exp-root", default="exp")
    ap.add_argument("--out-dir",  default="exp/summary_0920")
    ap.add_argument("--datasets", nargs="*", default=["scannet","nuscenes","s3dis","modelnet"])
    args = ap.parse_args()

    exp_root = Path(args.exp_root)
    out_dir  = Path(args.out_dir); out_dir.mkdir(parents=True, exist_ok=True)

    rows=[]
    for run_dir in sorted([p for p in exp_root.iterdir() if p.is_dir()]):
        ds = guess_dataset(run_dir)
        if ds not in args.datasets: continue
        logs = find_logs(run_dir)
        if not logs: continue

        w,a,mode = parse_bits_and_mode_by_name_or_log(run_dir)
        ckpt = find_ckpt(run_dir)
        avg=qratio=None
        if ckpt is not None:
            sd = load_state_dict_any(ckpt)
            if sd:
                force_fp32 = (mode=="fp32")
                ab, qr = avg_weight_bits(sd, w_bits=w, force_fp32=force_fp32)
                if not math.isnan(ab):
                    avg = round(ab,3); qratio = round(qr,2)

        kind = DATASET_KIND.get(ds, "seg")
        if kind=="seg":
            score = parse_miou_from_logs(logs); key="mIoU_best"
        else:
            score = parse_acc_from_logs(logs);  key="Acc_best"

        gflops = parse_flops_from_logs(logs)
        ms_it, fps = parse_speed_from_logs(logs)

        rows.append(dict(
            dataset=ds, run=run_dir.name, mode=mode, w_bits=w, a_bits=a,
            avg_weight_bit=avg, quant_ratio=qratio,
            FLOPs_GFLOPs=gflops, ms_per_it=ms_it, fps=fps,
            **{key: score},
            ckpt=str(ckpt) if ckpt else ""
        ))

    # CSV
    csv_path = out_dir/"summary_all_0920.csv"
    with csv_path.open("w", newline="") as f:
        cols=["dataset","run","mode","w_bits","a_bits","avg_weight_bit","quant_ratio",
              "mIoU_best","Acc_best","FLOPs_GFLOPs","ms_per_it","fps","ckpt"]
        w = csv.DictWriter(f, fieldnames=cols); w.writeheader()
        for r in rows: w.writerow({k: r.get(k,"") for k in cols})
    print(f"[OK] CSV: {csv_path}")

    # 画单张卡片 + 对比卡片
    by_ds = {}  # 修复：从空字典开始，避免 setdefault 不生效
    for r in rows:
        ds = r["dataset"]
        kind = DATASET_KIND.get(ds, "seg")
        if ds not in by_ds:
            by_ds[ds] = {"kind": kind, "base": [], "quant": []}
        else:
            # 再保险：确保键都在
            by_ds[ds].setdefault("kind", kind)
            by_ds[ds].setdefault("base", [])
            by_ds[ds].setdefault("quant", [])

        # 判断 base/quant
        if r.get("avg_weight_bit") is not None:
            tag = "base" if r["avg_weight_bit"] >= 31.9 else "quant"
        else:
            tag = "fp32" if r["mode"] == "fp32" else ("quant" if r["mode"] == "quant" else None)
            tag = "base" if tag == "fp32" else tag

        if tag in ("base", "quant"):
            by_ds[ds][tag].append(r)

    def score_key(kind): return "mIoU_best" if kind=="seg" else "Acc_best"

    for ds, bucket in by_ds.items():
        kind = bucket.get("kind","seg")
        key  = score_key(kind)
        # best 选择策略：优先有分数的最大值
        for tag in ["base","quant"]:
            cands=bucket.get(tag,[])
            if not cands:
                print(f"[skip] {ds} {tag}: 没有可用条目"); continue
            got=[x for x in cands if x.get(key) is not None]
            best=sorted(got, key=lambda x: x[key], reverse=True)[0] if got else cands[0]
            info={"avg_bit":best.get("avg_weight_bit"),"qratio":best.get("quant_ratio"),
                  "flops":best.get("FLOPs_GFLOPs"),"ms_per_it":best.get("ms_per_it"),"fps":best.get("fps")}
            if kind=="seg": info["miou"]=best.get("mIoU_best")
            else:           info["acc"]=best.get("Acc_best")
            out_png = out_dir / f"{ds}_{tag}_0920.png"
            draw_card_png(ds, tag, kind, info, str(out_png))
            bucket[f"{tag}_best"]=info

        # 对比卡片
        if "base_best" in bucket and "quant_best" in bucket:
            out_cmp = out_dir / f"{ds}_compare_0920.png"
            draw_compare_png(ds, kind, bucket["base_best"], bucket["quant_best"], str(out_cmp))
        else:
            print(f"[skip] {ds} compare: base 或 quant 缺失")

if __name__ == "__main__":
    main()