"""多分辨率网格生成 —— 端到端推理。

  低分辨率生成器(TSSR 三前向离散扩散) 逐级生成到格 64, 每级 20 步;
  res32 / res64 两级产物各补一次洞再喂下一级;
  高分辨率细分器继续 64 -> 128 -> 256 -> 512, 每级 20 步, 全程补洞。

用法:
  python scripts/infer.py --data <mesh 目录> --out_dir outputs/run1
  python scripts/infer.py --data <mesh 目录> --out_dir outputs/run1 --ncase 20 --empty_bias 0.3
"""
import argparse
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

DEFAULTS = dict(
    STAGE="2,4,8,16,32,64",                                  # 生成器级联阶梯(产物真实体素格)
    STEPS="20",                                              # 生成器每级扩散步数
    SUPERRES="1", SUPERRES_TO="512", SUPERRES_STEPS="20",    # 细分器 64->512, 每级步数
    FILL="1",                                                # 级联内补洞(res32/res64, 见 FILL_MIN_RES)
    SR_FILL_IN="1", SUPERRES_FIX_HOLES="1",                  # 细分入口 + 每级输出补洞
    DECIMATE_RATIO="0.5",                                    # 输入预塌缩比(0=不塌缩)
    METHODS="plain", MULTI_PC="0",
    SAVE_GT_LEVELS="1", SAVE_PC_LEVELS="1",
    CONFIG=os.path.join(ROOT, "configs", "generator.yaml"),
    SUPERRES_CONFIG=os.path.join(ROOT, "configs", "subdivider.yaml"),
)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--data", required=True, help="输入 mesh 目录(glb/obj/ply 混合)")
    ap.add_argument("--out_dir", default=os.path.join(ROOT, "outputs"), help="输出根目录")
    ap.add_argument("--ckpt", default=os.path.join(ROOT, "ckpt", "generator.pth"), help="生成器权重")
    ap.add_argument("--subdivider_ckpt", default=os.path.join(ROOT, "ckpt", "subdivider.pth"), help="细分器权重")
    ap.add_argument("--out_suffix", default="", help="输出子目录后缀")
    ap.add_argument("--ncase", type=int, default=0, help="只跑前 N 个(文件名字母序); 0=全部")
    ap.add_argument("--case", default="", help="只跑指定 case, 逗号分隔文件名子串(建议带尾点: 'a.' 不会匹配 'a1.glb')")
    ap.add_argument("--gpu", default=None, help="CUDA_VISIBLE_DEVICES")
    ap.add_argument("--steps", type=int, default=None, help="覆盖每级步数(默认 20)")
    ap.add_argument("--no_fill", action="store_true", help="关闭所有补洞")
    ap.add_argument("--fill_min_res", type=int, default=0, help=">0: 级联内补洞只在产物格>=该值的级做(默认 0=res32 起)")
    ap.add_argument("--empty_bias", type=float, default=0.0, help=">0: 连接头空类 logit 偏置, 稀疏化控面数(0.3 约 -25%% 面)")
    a = ap.parse_args()

    for k, v in DEFAULTS.items():
        os.environ.setdefault(k, v)
    os.environ["RAIL_CKPT"] = os.path.abspath(a.ckpt)
    os.environ["SUPERRES_CKPT"] = os.path.abspath(a.subdivider_ckpt)
    os.environ["BASE_DIR"] = os.path.dirname(os.path.abspath(a.ckpt))
    os.environ["OUT_ROOT"] = os.path.abspath(a.out_dir)
    os.environ["DATASET_DIR"] = os.path.abspath(a.data)
    os.environ.setdefault("DATASET", "custom")
    os.environ["NCASE"] = str(a.ncase if a.ncase > 0 else 10 ** 9)
    if a.case:          os.environ["CASE_STEM"] = a.case
    if a.out_suffix:    os.environ["OUT_SUFFIX"] = a.out_suffix
    if a.gpu:           os.environ["CUDA_VISIBLE_DEVICES"] = a.gpu
    if a.steps:         os.environ["STEPS"] = os.environ["SUPERRES_STEPS"] = str(a.steps)
    if a.no_fill:       os.environ.update(FILL="0", SR_FILL_IN="0", SUPERRES_FIX_HOLES="0")
    if a.fill_min_res:  os.environ["FILL_MIN_RES"] = str(a.fill_min_res)
    if a.empty_bias:    os.environ["EMPTY_BIAS"] = str(a.empty_bias)

    os.chdir(ROOT)
    import torch
    import modules.pipeline as P
    from models.generator import FactorizedTSSRTransformer

    class TSSRAdapter(FactorizedTSSRTransformer):
        """把主管线的 .sample() 转发到三前向 sample_tssr; forward 缺省补 mode=1(mask 步)。"""

        @torch.no_grad()
        def sample(self, coord_list, depth, pc=None, pc_target=None, T_infer=10, cfg_scale=1.0,
                   temperature=0.0, cfg_schedule="constant", cfg_start=1.0, **_ignored):
            return self.sample_tssr(coord_list, depth, pc=pc, pc_target=pc_target, T_infer=T_infer,
                                    cfg_scale=cfg_scale, cfg_schedule=cfg_schedule, cfg_start=cfg_start,
                                    mask_logits_lo=float(os.environ.get("MASK_LO", "0.2")),
                                    mask_logits_hi=float(os.environ.get("MASK_HI", "0.9")))

        def forward_varlen(self, *args, **kwargs):
            kwargs.pop("cond_cache", None)
            if kwargs.get("mode", None) is None and len(args) < 7:
                kwargs["mode"] = 1
            return super().forward_varlen(*args, **kwargs)

    P.BaseModel = TSSRAdapter
    P.main()


if __name__ == "__main__":
    main()