ptq_eval.py

# # eval.py — datasets: mnist/cifar10/cifar100/tinyimagenet/imagenet
# # Uses runner.py unified builders when available; keeps SQ + ACT wrap-before-load + DP key clean.

# import os, json, yaml, torch
# from tqdm import tqdm
# import torch.nn as nn
# from types import SimpleNamespace
# from collections import OrderedDict

# from utils.common import set_seed, accuracy
# from utils.checkpoint import load_checkpoint

# # ---------- Optional PTQ bits (import defensively) ----------
# wrap_input_quant = None
# apply_act_quant_from_files = None

# try:
#     from ptq.actquant import wrap_input_quant as _wrap_input_quant
#     wrap_input_quant = _wrap_input_quant
#     try:
#         from ptq.actquant import apply_act_quant_from_files as _apply
#         apply_act_quant_from_files = _apply
#     except Exception:
#         pass
# except Exception:
#     pass

# if wrap_input_quant is None or apply_act_quant_from_files is None:
#     try:
#         from quant.actquant import wrap_input_quant as _wrap_input_quant2
#         wrap_input_quant = wrap_input_quant or _wrap_input_quant2
#         try:
#             from quant.actquant import apply_act_quant_from_files as _apply2
#             apply_act_quant_from_files = apply_act_quant_from_files or _apply2
#         except Exception:
#             pass
#     except Exception:
#         pass

# try:
#     from ptq.uniform import wrap_input_uniform as wrap_input_uniform_act
# except Exception:
#     wrap_input_uniform_act = None

# try:
#     from ptq.smoothquant import load_scales_json as load_sq, apply_smoothquant
# except Exception:
#     load_sq = None
#     apply_smoothquant = None


# # ---------- Owners (can be skipped via cfg.ptq.skip_owner_actq) ----------
# OWNER_CLASSES = {
#     "GRAMLayer",
#     "KAGNConvNDLayer", "KAGNConv1DLayer", "KAGNConv2DLayer", "KAGNConv3DLayer",
# }


# # =====================================================================================
# # Try to REUSE runner.py unified builders (preferred)
# # =====================================================================================
# _HAVE_RUNNER = False
# try:
#     # runner.py has these helpers in your posted version
#     from runner import _adapt_config_to_args as _runner_adapt_config_to_args
#     from runner import _build_data as _runner_build_data
#     from runner import _build_model as _runner_build_model
#     _HAVE_RUNNER = True
# except Exception:
#     _HAVE_RUNNER = False


# # =====================================================================================
# # Fallback adapter (only used if runner helpers can’t be imported)
# # =====================================================================================
# def _adapt_config_to_args_fallback(cfg: dict):
#     """
#     Minimal clone of runner._adapt_config_to_args to support eval without importing runner.py.
#     """
#     try:
#         import munch
#         args = munch.munchify(cfg)
#         Munch = munch.Munch
#     except Exception:
#         # ultra-light fallback without munch
#         class Munch(dict):
#             __getattr__ = dict.get
#             __setattr__ = dict.__setitem__
#         def munchify(x):
#             if isinstance(x, dict):
#                 m = Munch()
#                 for k, v in x.items():
#                     m[k] = munchify(v)
#                 return m
#             return x
#         args = munchify(cfg)
#         Munch = Munch

#     if not getattr(args, "dataloader", None):
#         args.dataloader = Munch()

#     dl = args.dataloader

#     # dataset
#     if getattr(dl, "dataset", None) is None:
#         dl.dataset = (cfg.get("data", {}).get("dataset")
#                       or cfg.get("dataloader", {}).get("dataset")
#                       or "cifar10")

#     # path (new configs)
#     if getattr(dl, "path", None) is None:
#         dl.path = (cfg.get("data", {}).get("root")
#                    or cfg.get("dataloader", {}).get("path")
#                    or "./datasets")

#     # batch_size
#     if getattr(dl, "batch_size", None) is None:
#         dl.batch_size = (cfg.get("data", {}).get("batch_size")
#                          or cfg.get("dataloader", {}).get("batch_size")
#                          or 128)

#     # val_split
#     if getattr(dl, "val_split", None) is None:
#         dl.val_split = (cfg.get("data", {}).get("val_split")
#                         or cfg.get("dataloader", {}).get("val_split")
#                         or 0.0)

#     # workers
#     if getattr(dl, "workers", None) is None:
#         dl.workers = (cfg.get("data", {}).get("num_workers")
#                       or cfg.get("dataloader", {}).get("workers")
#                       or 4)

#     # deterministic
#     if getattr(dl, "deterministic", None) is None:
#         dl.deterministic = bool(cfg.get("deterministic", True))

#     # num_classes
#     if getattr(dl, "num_classes", None) is None:
#         dl.num_classes = (cfg.get("model", {}).get("num_classes")
#                           or cfg.get("dataloader", {}).get("num_classes")
#                           or cfg.get("num_classes")
#                           or 10)

#     # arch (top-level for create_model)
#     if getattr(args, "arch", None) is None:
#         args.arch = (cfg.get("model", {}).get("arch")
#                      or cfg.get("arch")
#                      or "kan_test_all_mnist")

#     # pre_trained
#     if getattr(args, "pre_trained", None) is None:
#         args.pre_trained = bool(cfg.get("pretrained", {}).get("pretrain_from_fp32", False))

#     return args


# def _resolve_seed(cfg: dict) -> int:
#     return int(cfg.get("seed", 42))


# def _resolve_dataset(cfg: dict) -> str:
#     return str(
#         cfg.get("data", {}).get("dataset")
#         or cfg.get("dataloader", {}).get("dataset", "cifar10")
#     ).lower()


# def _resolve_device(cfg: dict, override: str | None = None) -> str:
#     """
#     Supports:
#       - old: device: "cuda"
#       - new: device: {type: cuda, gpu: [0]}
#     """
#     if override:
#         dev = override
#     else:
#         dev_cfg = cfg.get("device", "cuda")
#         if isinstance(dev_cfg, dict):
#             dev = str(dev_cfg.get("type", "cuda"))
#             gpus = dev_cfg.get("gpu", None)
#             if dev.startswith("cuda") and isinstance(gpus, (list, tuple)) and len(gpus) > 0:
#                 dev = f"cuda:{int(gpus[0])}"
#         else:
#             dev = str(dev_cfg)

#     if "cuda" in dev and not torch.cuda.is_available():
#         return "cpu"
#     return dev


# def _clean_state_dict(raw_sd):
#     """
#     Accepts either a full checkpoint dict or a bare state_dict.
#     Strips common DP prefixes: module./model./net/.
#     """
#     if isinstance(raw_sd, dict) and any(k in raw_sd for k in ["state_dict", "model", "model_state", "weights"]):
#         sd = raw_sd.get("state_dict") or raw_sd.get("model") or raw_sd.get("model_state") or raw_sd.get("weights")
#     else:
#         sd = raw_sd
#     if not isinstance(sd, dict):
#         raise RuntimeError("Checkpoint format not recognized: expected a dict or a dict with 'state_dict'.")
#     cleaned = OrderedDict()
#     for k, v in sd.items():
#         if k.startswith("module."):
#             k = k[len("module."):]
#         elif k.startswith("model."):
#             k = k[len("model."):]
#         elif k.startswith("net."):
#             k = k[len("net."):]
#         cleaned[k] = v
#     return cleaned


# def _needs_act_wrap_from_keys(state_dict_keys) -> bool:
#     # your heuristic
#     for k in state_dict_keys:
#         if k.endswith(".scale") or ".layer.weight" in k or ".layer.bias" in k:
#             return True
#     return False


# def _load_json_safe(path: str):
#     try:
#         with open(path, "r") as f:
#             return json.load(f)
#     except Exception:
#         return None


# def _is_uniform_actseq(mod: nn.Module) -> bool:
#     return isinstance(mod, nn.Sequential) and len(mod) == 2 \
#         and (mod[0].__class__.__name__ == "UniformActQuant") \
#         and isinstance(mod[1], (nn.Conv2d, nn.Linear))


# def _is_act_wrapper(mod: nn.Module) -> bool:
#     return mod.__class__.__name__ in {
#         "InputQuantWrapper",
#         "KANLinearActQuant",
#         "KANLayerActQuant",
#         "GRAMActQuant",
#         "KAGNConvActQuant",
#     }


# def _remap_state_dict_for_uniform_wrappers(model: nn.Module, sd: dict) -> dict:
#     """
#     For Sequential(UniformActQuant, Conv/Linear):
#       '<name>.weight' -> '<name>.1.weight'
#       '<name>.bias'   -> '<name>.1.bias'
#     Also seed '<name>.0.s' and '<name>.0.zp' from model buffers if missing.
#     """
#     name_to_mod = dict(model.named_modules())
#     wrapped = [nm for nm, m in name_to_mod.items() if _is_uniform_actseq(m)]
#     if not wrapped:
#         return sd
#     sd = dict(sd)
#     model_sd = model.state_dict()
#     for nm in wrapped:
#         for buf in (f"{nm}.0.s", f"{nm}.0.zp"):
#             if buf not in sd and buf in model_sd:
#                 sd[buf] = model_sd[buf]
#         fw, fb = f"{nm}.weight", f"{nm}.bias"
#         w1, b1 = f"{nm}.1.weight", f"{nm}.1.bias"
#         if (fw in sd) and (w1 not in sd):
#             sd[w1] = sd.pop(fw)
#         if (fb in sd) and (b1 not in sd):
#             sd[b1] = sd.pop(fb)
#     return sd


# def _remap_state_dict_for_act_wrappers(model: nn.Module, sd: dict) -> dict:
#     """
#     After inserting ACT wrappers, model expects '.layer.' between wrapper root and children:
#       'root.weight'  -->  'root.layer.weight'
#     """
#     sd = dict(sd)
#     name_to_mod = dict(model.named_modules())
#     wrap_roots = sorted([nm for nm, m in name_to_mod.items() if _is_act_wrapper(m)], key=len, reverse=True)
#     if not wrap_roots:
#         return sd
#     new_sd = {}
#     for k, v in sd.items():
#         remapped = False
#         for root in wrap_roots:
#             pref = root + "."
#             layer_pref = root + ".layer."
#             if k.startswith(layer_pref):
#                 new_sd[k] = v
#                 remapped = True
#                 break
#             if k.startswith(pref):
#                 new_sd[layer_pref + k[len(pref):]] = v
#                 remapped = True
#                 break
#         if not remapped:
#             new_sd[k] = v
#     return new_sd


# def _print_first_wrapper(model: nn.Module):
#     for n, m in model.named_modules():
#         if hasattr(m, "abit"):
#             print("FIRST WRAPPER:", n, "abit=", getattr(m, "abit"))
#             break


# def _apply_actq_sidecars_before_load(
#     model: nn.Module,
#     act_scales_json: str,
#     act_params_json: str,
#     cfg: dict,
#     layer_types_act: tuple,
# ):
#     """
#     Prefer params+scales (correct abit/signed) if available.
#     Falls back to legacy behavior (cfg.ptq.a_bit_default) if params missing.
#     """
#     if not os.path.exists(act_scales_json):
#         return model, False, False  # (model, wrapped, used_uniform)

#     # If params exist, use them (and respect skip_owner via layer_types_act)
#     if os.path.exists(act_params_json) and wrap_input_quant is not None:
#         params = _load_json_safe(act_params_json) or {}
#         scales = _load_json_safe(act_scales_json) or {}
#         abit = int(params.get("abit", cfg.get("ptq", {}).get("a_bit_default", 8)))
#         signed = bool(params.get("signed", True))
#         model = wrap_input_quant(model, scales, abit=abit, layer_types=layer_types_act, signed=signed)
#         print(f"[eval] Activation wrappers inserted from {act_scales_json} + {act_params_json}"
#               f" (abit={abit}, signed={signed})"
#               f"{' (owners skipped)' if bool(cfg.get('ptq', {}).get('skip_owner_actq', False)) else ''}.")
#         _print_first_wrapper(model)
#         return model, True, False

#     # Scales-only legacy path
#     if wrap_input_quant is not None:
#         try:
#             act_scales = _load_json_safe(act_scales_json) or {}
#             abit = int(cfg.get("ptq", {}).get("a_bit_default", 8))
#             model = wrap_input_quant(model, act_scales, abit=abit, layer_types=layer_types_act)
#             print(f"[eval] Activation wrappers inserted from {act_scales_json}"
#                   f" (NO params; abit default={abit})"
#                   f"{' (owners skipped)' if bool(cfg.get('ptq', {}).get('skip_owner_actq', False)) else ''}.")
#             _print_first_wrapper(model)
#             return model, True, False
#         except Exception as e:
#             print(f"[eval] Warning: legacy wrap_input_quant failed: {e}")

#     # Optional uniform fallback
#     if wrap_input_uniform_act is not None:
#         params = _load_json_safe(act_scales_json)
#         if isinstance(params, dict) and params:
#             sample = next(iter(params.values())) if len(params) else {}
#             if isinstance(sample, dict) and "meta" in sample:
#                 try:
#                     model = wrap_input_uniform_act(model, params, layer_types=("Conv2d","Linear"))
#                     print("[eval] Activation wrappers inserted via ptq.uniform before weight load.")
#                     return model, True, True
#                 except Exception as e:
#                     print(f"[eval] Warning: ptq.uniform wrap failed: {e}")

#     return model, False, False


# # =====================================================================================
# # Unified builders (prefer runner.py, else fallback)
# # =====================================================================================
# def _build_test_loader(cfg: dict):
#     """
#     Prefer runner/util unified loader so eval matches training pipelines.
#     Returns test_loader.
#     """
#     if _HAVE_RUNNER:
#         # runner._build_data returns (train_loader, test_loader)
#         _, te = _runner_build_data(cfg)
#         return te

#     # Fallback: call util.load_data(args.dataloader) directly
#     import util
#     args = _adapt_config_to_args_fallback(cfg)

#     # Allow test_batch_size override if present in old cfg.data.test_batch_size
#     # (util.load_data only sees args.dataloader.batch_size)
#     test_bs = cfg.get("data", {}).get("test_batch_size", None)
#     if test_bs is not None:
#         args.dataloader.batch_size = int(test_bs)

#     tr, va, te = util.load_data(args.dataloader)
#     return te


# def _build_model(cfg: dict):
#     if _HAVE_RUNNER:
#         return _runner_build_model(cfg)
#     from model import create_model
#     args = _adapt_config_to_args_fallback(cfg)
#     return create_model(args)


# # =====================================================================================
# # Main eval
# # =====================================================================================
# @torch.no_grad()
# def evaluate_ckpt(config_path: str, ckpt_path: str, device: str | None = None) -> float:
#     cfg = yaml.safe_load(open(config_path, "r"))

#     seed = _resolve_seed(cfg)
#     det = bool(cfg.get("deterministic", True)) or bool(cfg.get("dataloader", {}).get("deterministic", True))
#     set_seed(seed, deterministic=det)

#     device = _resolve_device(cfg, override=device)

#     # Build test loader + model skeleton using unified pipeline
#     te_loader = _build_test_loader(cfg)
#     model = _build_model(cfg).to(device)

#     # Checkpoint + sidecars
#     ckpt = load_checkpoint(ckpt_path, map_location="cpu")
#     sd = _clean_state_dict(ckpt)

#     base = os.path.splitext(ckpt_path)[0]
#     act_scales_json = base + "_act_scales.json"
#     act_params_json = base + "_act_params.json"   # optional but preferred
#     smooth_json     = base + "_smooth_scales.json"

#     # Layer type lists (for SQ + ACT)
#     layer_types_all = tuple(cfg.get("ptq", {}).get("layer_types", (
#         "Conv2d","Linear","SplineLinear","KANLinear","KANLayer",
#         "GRAMLayer","KAGNConvNDLayer","KAGNConv1DLayer","KAGNConv2DLayer","KAGNConv3DLayer"
#     )))
#     skip_owner = bool(cfg.get("ptq", {}).get("skip_owner_actq", False))
#     layer_types_act = tuple(t for t in layer_types_all if (not skip_owner) or (t not in OWNER_CLASSES))

#     wrapped = False
#     used_uniform = False

#     # 1) SmoothQuant BEFORE weight load
#     if os.path.exists(smooth_json) and load_sq is not None and apply_smoothquant is not None:
#         try:
#             sq_scales = load_sq(smooth_json)
#             apply_smoothquant(model, sq_scales, layer_types=layer_types_all)
#             print(f"[eval] Applied SmoothQuant scales: {smooth_json}")
#         except Exception as e:
#             print(f"[eval] SmoothQuant scales not applied: {e}")

#     # 2) Activation wrappers BEFORE weight load (prefer params json)
#     model, wrapped, used_uniform = _apply_actq_sidecars_before_load(
#         model=model,
#         act_scales_json=act_scales_json,
#         act_params_json=act_params_json,
#         cfg=cfg,
#         layer_types_act=layer_types_act,
#     )
#     model = model.to(device)

#     # Heuristic: checkpoint expects wrappers even without JSON
#     if (not wrapped) and _needs_act_wrap_from_keys(sd.keys()) and wrap_input_quant is not None:
#         try:
#             abit = int(cfg.get("ptq", {}).get("a_bit_default", 8))
#             model = wrap_input_quant(model, {}, abit=abit, layer_types=layer_types_act).to(device)
#             print(f"[eval] Detected wrapped checkpoint; created wrapper structure before load"
#                   f"{' (owners skipped)' if skip_owner else ''}.")
#             _print_first_wrapper(model)
#             wrapped = True
#         except Exception as e:
#             print(f"[eval] Warning: could not create wrapper structure: {e}")

#     # 3) Remap state_dict if wrappers changed structure
#     if wrapped:
#         if used_uniform:
#             sd = _remap_state_dict_for_uniform_wrappers(model, sd)
#         sd = _remap_state_dict_for_act_wrappers(model, sd)

#     # 4) Load weights
#     try:
#         model.load_state_dict(sd, strict=True)
#     except Exception as e:
#         print(f"[eval] strict=True failed: {e}\n[eval] retrying with strict=False")
#         model.load_state_dict(sd, strict=False)

#     # 5) Eval
#     model.eval()
#     top1_sum, n = 0.0, 0
#     for images, targets in tqdm(te_loader, desc="Eval"):
#         images, targets = images.to(device), targets.to(device)
#         logits = model(images)
#         acc1 = accuracy(logits, targets, topk=(1,))[0].item()  # typically in %
#         top1_sum += acc1 * images.size(0)
#         n += images.size(0)

#     return top1_sum / max(1, n)


# if __name__ == "__main__":
#     import argparse
#     ap = argparse.ArgumentParser()
#     ap.add_argument("--config", required=True)
#     ap.add_argument("--ckpt", required=True)
#     ap.add_argument("--device", type=str, default=None)
#     args = ap.parse_args()
#     acc = evaluate_ckpt(args.config, args.ckpt, device=args.device)
#     print(f"Top-1 Acc: {acc:.2f}")



# eval.py — datasets: mnist/cifar10/cifar100/tinyimagenet/imagenet
# Uses runner.py unified builders when available; keeps SQ + ACT wrap-before-load + DP key clean
# + runner-style logging (console + optional file)

import os, json, yaml, torch
from tqdm import tqdm
import torch.nn as nn
from collections import OrderedDict
from types import SimpleNamespace

from util.ptq_helpers import set_seed, accuracy, load_checkpoint


# =============================================================================
# Optional PTQ bits (import defensively)
# =============================================================================
wrap_input_quant = None
apply_act_quant_from_files = None

try:
    from ptq.actquant import wrap_input_quant as _wrap_input_quant
    wrap_input_quant = _wrap_input_quant
    try:
        from ptq.actquant import apply_act_quant_from_files as _apply
        apply_act_quant_from_files = _apply
    except Exception:
        pass
except Exception:
    pass

if wrap_input_quant is None or apply_act_quant_from_files is None:
    try:
        from quant.actquant import wrap_input_quant as _wrap_input_quant2
        wrap_input_quant = wrap_input_quant or _wrap_input_quant2
        try:
            from quant.actquant import apply_act_quant_from_files as _apply2
            apply_act_quant_from_files = apply_act_quant_from_files or _apply2
        except Exception:
            pass
    except Exception:
        pass

try:
    from ptq.uniform import wrap_input_uniform as wrap_input_uniform_act
except Exception:
    wrap_input_uniform_act = None

try:
    from ptq.smoothquant import load_scales_json as load_sq, apply_smoothquant
except Exception:
    load_sq = None
    apply_smoothquant = None


# =============================================================================
# Runner-like logging utilities (copied pattern)
# =============================================================================
import logging
from pathlib import Path
import time

def setup_logging(
    log_dir: Path = None,
    experiment_name: str = "EVAL",
    level: int = logging.INFO,
):
    logger = logging.getLogger()

    # Clear any existing handlers
    logger.handlers.clear()
    logger.setLevel(level)

    formatter = logging.Formatter(
        fmt="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    )

    # Console handler
    console_handler = logging.StreamHandler()
    console_handler.setLevel(level)
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)

    # File handler (optional)
    if log_dir is not None:
        log_dir = Path(log_dir)
        log_dir.mkdir(parents=True, exist_ok=True)

        time_str = time.strftime("%Y%m%d-%H%M%S")
        log_filename = f"{experiment_name}_{time_str}.log"
        log_file = log_dir / log_filename

        file_handler = logging.FileHandler(log_file, mode="w")
        file_handler.setLevel(level)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

        logger.info(f"Log file: {log_file}")

    return logger, log_dir

def get_logger():
    return logging.getLogger()


# =============================================================================
# Owners (can be skipped via cfg.ptq.skip_owner_actq)
# =============================================================================
OWNER_CLASSES = {
    "GRAMLayer",
    "KAGNConvNDLayer", "KAGNConv1DLayer", "KAGNConv2DLayer", "KAGNConv3DLayer",
}


# =============================================================================
# Try to reuse runner.py unified builders
# =============================================================================
_HAVE_RUNNER = False
try:
    from runner import _build_data as _runner_build_data
    from runner import _build_model as _runner_build_model
    _HAVE_RUNNER = True
except Exception:
    _HAVE_RUNNER = False


# =============================================================================
# Config helpers
# =============================================================================
def _resolve_seed(cfg: dict) -> int:
    return int(cfg.get("seed", 42))

def _resolve_dataset(cfg: dict) -> str:
    return str(
        cfg.get("data", {}).get("dataset")
        or cfg.get("dataloader", {}).get("dataset", "cifar10")
    ).lower()

def _resolve_device(cfg: dict, override: str | None = None) -> str:
    """
    Supports:
      - old: device: "cuda"
      - new: device: {type: cuda, gpu: [0]}
    """
    if override:
        dev = override
    else:
        dev_cfg = cfg.get("device", "cuda")
        if isinstance(dev_cfg, dict):
            dev = str(dev_cfg.get("type", "cuda"))
            gpus = dev_cfg.get("gpu", None)
            if dev.startswith("cuda") and isinstance(gpus, (list, tuple)) and len(gpus) > 0:
                dev = f"cuda:{int(gpus[0])}"
        else:
            dev = str(dev_cfg)

    if "cuda" in dev and not torch.cuda.is_available():
        return "cpu"
    return dev


def _clean_state_dict(raw_sd):
    """
    Accepts either a full checkpoint dict or a bare state_dict.
    Strips common DP prefixes: module./model./net/.
    """
    if isinstance(raw_sd, dict) and any(k in raw_sd for k in ["state_dict", "model", "model_state", "weights"]):
        sd = raw_sd.get("state_dict") or raw_sd.get("model") or raw_sd.get("model_state") or raw_sd.get("weights")
    else:
        sd = raw_sd
    if not isinstance(sd, dict):
        raise RuntimeError("Checkpoint format not recognized: expected a dict or a dict with 'state_dict'.")
    cleaned = OrderedDict()
    for k, v in sd.items():
        if k.startswith("module."):
            k = k[len("module."):]
        elif k.startswith("model."):
            k = k[len("model."):]
        elif k.startswith("net."):
            k = k[len("net."):]
        cleaned[k] = v
    return cleaned


def _needs_act_wrap_from_keys(state_dict_keys) -> bool:
    for k in state_dict_keys:
        if k.endswith(".scale") or ".layer.weight" in k or ".layer.bias" in k:
            return True
    return False


def _load_json_safe(path: str):
    try:
        with open(path, "r") as f:
            return json.load(f)
    except Exception:
        return None


# =============================================================================
# Wrapper detection + remapping
# =============================================================================
def _is_uniform_actseq(mod: nn.Module) -> bool:
    return isinstance(mod, nn.Sequential) and len(mod) == 2 \
        and (mod[0].__class__.__name__ == "UniformActQuant") \
        and isinstance(mod[1], (nn.Conv2d, nn.Linear))

def _is_act_wrapper(mod: nn.Module) -> bool:
    return mod.__class__.__name__ in {
        "InputQuantWrapper",
        "KANLinearActQuant",
        "KANLayerActQuant",
        "GRAMActQuant",
        "KAGNConvActQuant",
    }

def _remap_state_dict_for_uniform_wrappers(model: nn.Module, sd: dict) -> dict:
    name_to_mod = dict(model.named_modules())
    wrapped = [nm for nm, m in name_to_mod.items() if _is_uniform_actseq(m)]
    if not wrapped:
        return sd
    sd = dict(sd)
    model_sd = model.state_dict()
    for nm in wrapped:
        for buf in (f"{nm}.0.s", f"{nm}.0.zp"):
            if buf not in sd and buf in model_sd:
                sd[buf] = model_sd[buf]
        fw, fb = f"{nm}.weight", f"{nm}.bias"
        w1, b1 = f"{nm}.1.weight", f"{nm}.1.bias"
        if (fw in sd) and (w1 not in sd):
            sd[w1] = sd.pop(fw)
        if (fb in sd) and (b1 not in sd):
            sd[b1] = sd.pop(fb)
    return sd

def _remap_state_dict_for_act_wrappers(model: nn.Module, sd: dict) -> dict:
    sd = dict(sd)
    name_to_mod = dict(model.named_modules())
    wrap_roots = sorted([nm for nm, m in name_to_mod.items() if _is_act_wrapper(m)], key=len, reverse=True)
    if not wrap_roots:
        return sd
    new_sd = {}
    for k, v in sd.items():
        remapped = False
        for root in wrap_roots:
            pref = root + "."
            layer_pref = root + ".layer."
            if k.startswith(layer_pref):
                new_sd[k] = v
                remapped = True
                break
            if k.startswith(pref):
                new_sd[layer_pref + k[len(pref):]] = v
                remapped = True
                break
        if not remapped:
            new_sd[k] = v
    return new_sd

def _print_first_wrapper(model: nn.Module, logger):
    for n, m in model.named_modules():
        if hasattr(m, "abit"):
            logger.info(f"FIRST WRAPPER: {n} abit={getattr(m, 'abit')}")
            break


def _apply_actq_sidecars_before_load(
    model: nn.Module,
    act_scales_json: str,
    act_params_json: str,
    cfg: dict,
    layer_types_act: tuple,
    logger,
):
    if not os.path.exists(act_scales_json):
        return model, False, False

    # params+scales preferred
    if os.path.exists(act_params_json) and wrap_input_quant is not None:
        params = _load_json_safe(act_params_json) or {}
        scales = _load_json_safe(act_scales_json) or {}
        abit = int(params.get("abit", cfg.get("ptq", {}).get("a_bit_default", 8)))
        signed = bool(params.get("signed", True))
        model = wrap_input_quant(model, scales, abit=abit, layer_types=layer_types_act, signed=signed)
        logger.info(
            f"[eval] Activation wrappers inserted from {act_scales_json} + {act_params_json} "
            f"(abit={abit}, signed={signed})"
            f"{' (owners skipped)' if bool(cfg.get('ptq', {}).get('skip_owner_actq', False)) else ''}."
        )
        _print_first_wrapper(model, logger)
        return model, True, False

    # scales-only legacy path
    if wrap_input_quant is not None:
        try:
            act_scales = _load_json_safe(act_scales_json) or {}
            abit = int(cfg.get("ptq", {}).get("a_bit_default", 8))
            model = wrap_input_quant(model, act_scales, abit=abit, layer_types=layer_types_act)
            logger.info(
                f"[eval] Activation wrappers inserted from {act_scales_json} "
                f"(NO params; abit default={abit})"
                f"{' (owners skipped)' if bool(cfg.get('ptq', {}).get('skip_owner_actq', False)) else ''}."
            )
            _print_first_wrapper(model, logger)
            return model, True, False
        except Exception as e:
            logger.warning(f"[eval] legacy wrap_input_quant failed: {e}")

    # optional uniform fallback
    if wrap_input_uniform_act is not None:
        params = _load_json_safe(act_scales_json)
        if isinstance(params, dict) and params:
            sample = next(iter(params.values())) if len(params) else {}
            if isinstance(sample, dict) and "meta" in sample:
                try:
                    model = wrap_input_uniform_act(model, params, layer_types=("Conv2d", "Linear"))
                    logger.info("[eval] Activation wrappers inserted via ptq.uniform before weight load.")
                    return model, True, True
                except Exception as e:
                    logger.warning(f"[eval] ptq.uniform wrap failed: {e}")

    return model, False, False


# =============================================================================
# Unified builders
# =============================================================================
def _build_test_loader(cfg: dict, logger):
    if _HAVE_RUNNER:
        _, te = _runner_build_data(cfg)
        return te

    # Fallback: use util.load_data with a minimal dataloader object
    import util
    try:
        import munch
        args = munch.munchify(cfg)
        if not hasattr(args, "dataloader") or args.dataloader is None:
            args.dataloader = munch.Munch()
        dl = args.dataloader
    except Exception:
        class Obj(dict):
            __getattr__ = dict.get
            __setattr__ = dict.__setitem__
        args = Obj(cfg)
        if "dataloader" not in args:
            args["dataloader"] = Obj()
        dl = args["dataloader"]

    # Ensure required fields
    dl.dataset = getattr(dl, "dataset", None) or cfg.get("dataloader", {}).get("dataset", "cifar10")
    dl.path = getattr(dl, "path", None) or cfg.get("dataloader", {}).get("path") or cfg.get("data", {}).get("root") or "./datasets"
    dl.batch_size = int(getattr(dl, "batch_size", None) or cfg.get("data", {}).get("test_batch_size") or cfg.get("dataloader", {}).get("batch_size", 128))
    dl.val_split = float(getattr(dl, "val_split", None) or cfg.get("dataloader", {}).get("val_split", 0.0))
    dl.workers = int(getattr(dl, "workers", None) or cfg.get("data", {}).get("num_workers") or cfg.get("dataloader", {}).get("workers", 4))
    dl.deterministic = bool(getattr(dl, "deterministic", None) if getattr(dl, "deterministic", None) is not None else cfg.get("deterministic", True))

    tr, va, te = util.load_data(dl)
    return te


def _build_model(cfg: dict, logger):
    if _HAVE_RUNNER:
        return _runner_build_model(cfg)

    from model import create_model
    # minimal args
    dataset = _resolve_dataset(cfg)
    arch = str(cfg.get("model", {}).get("arch") or cfg.get("arch", "kagn"))
    args = SimpleNamespace(
        arch=arch,
        dataloader=SimpleNamespace(dataset=dataset),
        pre_trained=bool(cfg.get("pretrained", {}).get("pretrain_from_fp32", False)),
    )
    return create_model(args)


# =============================================================================
# Main eval
# =============================================================================
@torch.no_grad()
def evaluate_ckpt(config_path: str, ckpt_path: str, device: str | None = None, logger=None) -> float:
    logger = logger or get_logger()

    cfg = yaml.safe_load(open(config_path, "r"))

    seed = _resolve_seed(cfg)
    det = bool(cfg.get("deterministic", True)) or bool(cfg.get("dataloader", {}).get("deterministic", True))
    set_seed(seed, deterministic=det)

    device = _resolve_device(cfg, override=device)

    logger.info(f"Loaded config: {config_path}")
    logger.info(f"Checkpoint: {ckpt_path}")
    logger.info(f"Seed: {seed} | deterministic: {det}")
    logger.info(f"Device: {device}")
    logger.info(f"Runner builders available: {_HAVE_RUNNER}")

    # Data + model
    te_loader = _build_test_loader(cfg, logger)
    model = _build_model(cfg, logger).to(device)

    # Load checkpoint + sidecars
    ckpt = load_checkpoint(ckpt_path, map_location="cpu")
    sd = _clean_state_dict(ckpt)

    base = os.path.splitext(ckpt_path)[0]
    act_scales_json = base + "_act_scales.json"
    act_params_json = base + "_act_params.json"
    smooth_json     = base + "_smooth_scales.json"

    layer_types_all = tuple(cfg.get("ptq", {}).get("layer_types", (
        "Conv2d","Linear","SplineLinear","KANLinear","KANLayer",
        "GRAMLayer","KAGNConvNDLayer","KAGNConv1DLayer","KAGNConv2DLayer","KAGNConv3DLayer"
    )))
    skip_owner = bool(cfg.get("ptq", {}).get("skip_owner_actq", False))
    layer_types_act = tuple(t for t in layer_types_all if (not skip_owner) or (t not in OWNER_CLASSES))

    wrapped = False
    used_uniform = False

    # 1) SmoothQuant BEFORE load
    if os.path.exists(smooth_json) and load_sq is not None and apply_smoothquant is not None:
        try:
            sq_scales = load_sq(smooth_json)
            apply_smoothquant(model, sq_scales, layer_types=layer_types_all)
            logger.info(f"[eval] Applied SmoothQuant scales: {smooth_json}")
        except Exception as e:
            logger.warning(f"[eval] SmoothQuant scales not applied: {e}")

    # 2) ACT wrappers BEFORE load
    model, wrapped, used_uniform = _apply_actq_sidecars_before_load(
        model=model,
        act_scales_json=act_scales_json,
        act_params_json=act_params_json,
        cfg=cfg,
        layer_types_act=layer_types_act,
        logger=logger,
    )
    model = model.to(device)

    # Heuristic: checkpoint expects wrappers even without JSON
    if (not wrapped) and _needs_act_wrap_from_keys(sd.keys()) and wrap_input_quant is not None:
        try:
            abit = int(cfg.get("ptq", {}).get("a_bit_default", 8))
            model = wrap_input_quant(model, {}, abit=abit, layer_types=layer_types_act).to(device)
            logger.info(
                f"[eval] Detected wrapped checkpoint; created wrapper structure before load"
                f"{' (owners skipped)' if skip_owner else ''}."
            )
            _print_first_wrapper(model, logger)
            wrapped = True
        except Exception as e:
            logger.warning(f"[eval] could not create wrapper structure: {e}")

    # 3) Remap state_dict if structure changed
    if wrapped:
        if used_uniform:
            sd = _remap_state_dict_for_uniform_wrappers(model, sd)
        sd = _remap_state_dict_for_act_wrappers(model, sd)

    # 4) Load weights
    try:
        model.load_state_dict(sd, strict=True)
        logger.info("[eval] Loaded state_dict (strict=True)")
    except Exception as e:
        logger.warning(f"[eval] strict=True failed: {e} | retrying strict=False")
        model.load_state_dict(sd, strict=False)
        logger.info("[eval] Loaded state_dict (strict=False)")

    # 5) Evaluate
    model.eval()
    top1_sum, n = 0.0, 0

    logger.info("Starting evaluation loop...")
    for images, targets in tqdm(te_loader, desc="Eval"):
        images, targets = images.to(device), targets.to(device)
        logits = model(images)
        acc1 = accuracy(logits, targets, topk=(1,))[0].item()  # typically in %
        top1_sum += acc1 * images.size(0)
        n += images.size(0)

    top1 = top1_sum / max(1, n)
    logger.info(f"Done. Top-1 Acc: {top1:.2f}")
    return top1

from pathlib import Path

def _ckpt_stem(path: str) -> str:
    name = Path(path).name if path else "no_ckpt"
    for suf in (".pth.tar", ".pt", ".pth", ".bin"):
        if name.endswith(suf):
            return name[: -len(suf)]
    return Path(name).stem
# =============================================================================
# CLI
# =============================================================================
def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--device", type=str, default=None)

    # runner-like logging flags
    ap.add_argument("--log_dir", type=str, default=None,
                    help="Directory to save log files (default: derived from ckpt dir)")
    ap.add_argument("--log_level", type=str, default="INFO",
                    choices=["DEBUG", "INFO", "WARNING", "ERROR"],
                    help="Logging level (default: INFO)")
    args = ap.parse_args()

    # Determine experiment name from config (runner-like)
    cfg = yaml.safe_load(open(args.config, "r"))
    # experiment_name = cfg.get("name", "EVAL")
    cfg_name = cfg.get("name", "eval")
    ckpt_name = _ckpt_stem(getattr(args, "ckpt", ""))
    experiment_name = f"EVAL-{cfg_name}_{ckpt_name}"
    

    # Determine log dir (runner-like)
    if args.log_dir:
        log_dir = Path(args.log_dir)
    else:
        # put logs beside ckpt by default
        log_dir = Path(args.ckpt).resolve().parent / "logs"

    level = getattr(logging, args.log_level.upper(), logging.INFO)
    logger, _ = setup_logging(log_dir=log_dir, experiment_name=experiment_name, level=level)

    logger.info("=" * 70)
    logger.info("Starting EVAL")
    logger.info("=" * 70)

    acc = evaluate_ckpt(args.config, args.ckpt, device=args.device, logger=logger)
    print(f"Top-1 Acc: {acc:.2f}")  # keep simple print for scripts


if __name__ == "__main__":
    main()