import gc
import glob
import math
import os
import random
import re
from pathlib import Path

import numpy as np
import torch
import trimesh
from accelerate import Accelerator, DistributedDataParallelKwargs
from tqdm import tqdm

from utils.config import load_config
from dataset.mesh_dataset import mul_encode_mesh, discrete_undiscrete_and_clean
from models.subdivider import DiscreteDiffusionTransformer, sample_refine_tokens_with_topology, sample_refine_tokens_with_topology_fm
from modules.octree_codec_acc import MeshCodec
from utils.postprocessing import _fill_holes
from utils.data_utils import normalize
def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    os.environ["PYTHONHASHSEED"] = str(seed)

def _depth_from_resolution(resolution):
    if resolution <= 0:
        raise ValueError(f"resolution must be positive, got {resolution}")
    depth = int(round(math.log2(resolution)))
    if (1 << depth) != resolution:
        raise ValueError(f"resolution must be power of 2, got {resolution}")
    return depth


def restore_non_pad_token(token: torch.Tensor, shift: int, pad_id: int = 0):
    if shift == 0:
        return token
    restored = token.clone()
    non_pad_mask = restored != pad_id
    restored[non_pad_mask] = restored[non_pad_mask] + shift
    return restored


def read_dataset_paths(txt_path):
    mesh_paths = []
    with open(txt_path, "r") as f:
        for line in f:
            path = line.strip()
            if path:
                mesh_paths.append(path)
    return mesh_paths


def raw_stem_to_pair_key(raw_path):
    stem = Path(raw_path).stem
    # raw name example: 1(42)_y_f_4244 -> pair key: 1(42)
    return re.sub(r"_y_f_\d+$", "", stem)


def gt_stem_to_pair_key(gt_path):
    return Path(gt_path).stem

def get_face_count(file_path):
    """
    使用正则表达式从文件名中提取面片数量
    支持格式: *_f{num}_v*.obj 或 *_facenum_{num}_*.glb
    """
    filename = Path(file_path).name
    match = re.search(r'_f(\d+)_v|_facenum_(\d+)', filename)
    if match:
        # group(1) 对应 _f(\d+)_v 里的数字
        # group(2) 对应 _facenum_(\d+) 里的数字
        return int(match.group(1) or match.group(2))
    return None

# def build_raw_gt_pairs(raw_txt_path, gt_txt_path):
#     raw_paths = read_dataset_paths(raw_txt_path)
#     gt_paths = read_dataset_paths(gt_txt_path)

#     # If two txt paths are the same source, pair each mesh with itself directly.
#     if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve():
#         return [(p, p) for p in raw_paths], []

#     gt_by_key = {}
#     for gt_path in gt_paths:
#         gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path

#     pairs = []
#     missing_gt = []
#     for raw_path in raw_paths:
#         key = raw_stem_to_pair_key(raw_path)
#         gt_path = gt_by_key.get(key)
#         if gt_path is None:
#             missing_gt.append(raw_path)
#             continue
#         pairs.append((raw_path, gt_path))

#     return pairs, missing_gt
def build_raw_gt_pairs(raw_txt_path, gt_txt_path, min_faces=None, max_faces=None):
    raw_paths = read_dataset_paths(raw_txt_path)
    gt_paths = read_dataset_paths(gt_txt_path)

    # ===== 新增过滤逻辑 =====
    filtered_raw_paths = []
    for p in raw_paths:
        if min_faces is None and max_faces is None:
            filtered_raw_paths.append(p)
            continue
            
        faces = get_face_count(p)
        if faces is not None:
            # 检查是否在指定范围内
            if min_faces is not None and faces < min_faces:
                continue
            if max_faces is not None and faces > max_faces:
                continue
            filtered_raw_paths.append(p)
        else:
            # 如果文件名中没找到面片数特征，默认保留（你也可以改为 continue 丢弃它）
            filtered_raw_paths.append(p)
            
    raw_paths = filtered_raw_paths
    # ========================

    # If two txt paths are the same source, pair each mesh with itself directly.
    if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve():
        return [(p, p) for p in raw_paths], []

    gt_by_key = {}
    for gt_path in gt_paths:
        gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path

    pairs = []
    missing_gt = []
    for raw_path in raw_paths:
        key = raw_stem_to_pair_key(raw_path)
        gt_path = gt_by_key.get(key)
        if gt_path is None:
            missing_gt.append(raw_path)
            continue
        pairs.append((raw_path, gt_path))

    return pairs, missing_gt

def save_gt_mesh(mesh, out_dir, name):
    gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False)
    gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]]
    gt_mesh.export(os.path.join(out_dir, f"{name}_gt.obj"))

def save_raw_gt_mesh(mesh, out_dir, name):
    gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False)
    gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]]
    gt_mesh.export(os.path.join(out_dir, f"{name}_raw_gt.obj"))


def save_stage_input_mesh(mesh, out_dir, name, stage_in_res, stage_out_res, stage_in_depth):
    coarse_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False)
    coarse_mesh.vertices = coarse_mesh.vertices[:, [1, 2, 0]]
    suffix = f"_stage_{stage_in_res}to{stage_out_res}_input_coarse"
    coarse_mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{stage_in_depth}.obj"))


def save_stage_prediction(codec, raw_token, refine_token, out_dir, name, target_depth, suffix, fix_holes, raw_dis_mesh=None, gt_dis_mesh=None):
    eos_mask = raw_token[0] == 0
    if eos_mask.any():
        raw_tokens = raw_token[0, : int(torch.where(eos_mask)[0][0])].detach().cpu().long().tolist()
    else:
        raw_tokens = raw_token[0].detach().cpu().long().tolist()

    eos_mask_refine = refine_token[0] == 0
    if eos_mask_refine.any():
        refine_tokens = refine_token[0, : int(torch.where(eos_mask_refine)[0][0])].detach().cpu().long().tolist()
    else:
        refine_tokens = refine_token[0].detach().cpu().long().tolist()

    shared_len = min(len(raw_tokens), len(refine_tokens))
    block = (shared_len // 9) * 9
    if block <= 0:
        return None

    raw_tokens = raw_tokens[:block]
    refine_tokens = refine_tokens[:block]

    interleaved = []
    for i in range(0, block, 9):
        interleaved.extend(raw_tokens[i : i + 9])
        interleaved.extend(refine_tokens[i : i + 9])

    mesh = codec.decode_tokens(interleaved, target_depth=target_depth-1)
    if mesh is None:
        return None
    
    trimesh.repair.fix_normals(mesh)
    if fix_holes:
        verts = torch.tensor(mesh.vertices, device="cuda", dtype=torch.float32)
        faces = torch.tensor(mesh.faces, device="cuda", dtype=torch.int32)
        new_verts, new_faces = _fill_holes(verts, faces, debug=False, verbose=False, only_fill_holes=True)
        mesh = trimesh.Trimesh(vertices=new_verts.cpu().numpy(), faces=new_faces.cpu().numpy())
    
    return_mesh = mesh.copy()
    mesh.vertices = mesh.vertices[:, [1, 2, 0]]
    mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{target_depth}.obj"))

    if raw_dis_mesh is not None:
        raw_dis_mesh.vertices = raw_dis_mesh.vertices[:, [1, 2, 0]]
        raw_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_raw_dis_depth{target_depth-1}.obj"))
    if gt_dis_mesh is not None:
        gt_dis_mesh.vertices = gt_dis_mesh.vertices[:, [1, 2, 0]]
        gt_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_gt_dis_depth{target_depth}.obj"))
    return return_mesh


def process_one_mesh(
    raw_mesh_path,
    gt_mesh_path,
    model,
    mask_token_id,
    codec,
    device,
    pad_id,
    token_shift,
    config,
    output_folder,
    stage_resolutions,
    fix_holes=True,
    use_gt_pc=True
):
    import time
    t_start = time.time()
    import time
    t_total_start = time.time()
    
    name = Path(raw_mesh_path).stem
    t_start = time.time()
    mesh = trimesh.load(raw_mesh_path, force="mesh")
    mesh.vertices = mesh.vertices[:, [2, 0, 1]]
    gt_mesh = trimesh.load(gt_mesh_path, force="mesh")
    gt_mesh.vertices = gt_mesh.vertices[:, [2, 0, 1]]
    normalize(gt_mesh)
    normalize(mesh)
    save_raw_gt_mesh(mesh, output_folder, name)
    save_gt_mesh(gt_mesh, output_folder, name)
    t_load = time.time()
    print(f"[{name}] Load & normalize time: {t_load - t_start:.4f}s")

    stage_depths = [_depth_from_resolution(x) for x in stage_resolutions]

    # First-stage input uses codec encode-decode result at the minimum resolution.
    min_stage_depth = stage_depths[0]
    t_start = time.time()
    init_raw_token, _, init_quantized_mesh = codec.encode_mesh(mesh, target_depth=min_stage_depth)
    init_mesh = codec.decode_raw_tokens(init_raw_token, target_depth=min_stage_depth)
    t_init = time.time()
    print(f"[{name}] Init encode & decode (depth={min_stage_depth}): {t_init - t_start:.4f}s")
    
    if init_mesh is not None:
        current_mesh = init_mesh
    elif init_quantized_mesh is not None:
        current_mesh = init_quantized_mesh
    else:
        current_mesh = mesh.copy()
        
    for stage_idx in range(len(stage_depths) - 1):
        stage_start_time = time.time()
        stage_in_depth = stage_depths[stage_idx]
        stage_out_depth = stage_depths[stage_idx + 1]
        stage_in_res = stage_resolutions[stage_idx]
        stage_out_res = stage_resolutions[stage_idx + 1]

        save_stage_input_mesh(
            mesh=current_mesh,
            out_dir=output_folder,
            name=name,
            stage_in_res=stage_in_res,
            stage_out_res=stage_out_res,
            stage_in_depth=stage_in_depth,
        )
        
        t_start = time.time()
        stage_data = mul_encode_mesh(
            raw_mesh=current_mesh,
            gt_mesh=gt_mesh,
            codec=codec,
            tokenizer_target_depth=stage_in_depth,
            pc_num=config.data_processing.pc_num,
            total_pc_num=config.data_processing.total_pc_num,
            use_gt_pc=use_gt_pc,
            clean=False
        )
        t_mul_encode = time.time()
        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | mul_encode_mesh time: {t_mul_encode - t_start:.4f}s")

        raw_token = stage_data["raw_token"].unsqueeze(0).to(device)
        pc = stage_data["pc"].unsqueeze(0).to(device)
    
        t_start = time.time()
        # refine_token = sample_refine_tokens_with_topology(
        #     model,
        #     mask_token_id,
        #     raw_token,
        #     pc,
        #     use_gumbel=False,
        #     gumbel_temp_init=1.0,
        #     use_soft_consensus=False,
        #     consensus_start_step_ratio=0,
        #     cfg_scale=1.0,
        #     T_infer=config.evaluation.diffusion_sample_steps,
        # )
        refine_token = sample_refine_tokens_with_topology_fm(
            model,
            mask_token_id,
            raw_token,
            pc,
            use_gumbel=False,
            gumbel_temp_init=1.0,
            use_soft_consensus=False,
            consensus_start_step_ratio=0,
            cfg_scale=1.0,
            T_infer=config.evaluation.diffusion_sample_steps,
        )
        
        torch.cuda.synchronize()
        t_sample = time.time()
        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | sample_refine_tokens_with_topology time: {t_sample - t_start:.4f}s")
        
        refine_token = restore_non_pad_token(refine_token, shift=token_shift, pad_id=pad_id)

        t_start = time.time()
        stage_suffix = f"_stage_{stage_in_res}to{stage_out_res}"
        stage_mesh = save_stage_prediction(
            codec=codec,
            raw_token=raw_token,
            refine_token=refine_token,
            out_dir=output_folder,
            name=name,
            target_depth=stage_out_depth,
            suffix=stage_suffix,
            fix_holes=fix_holes,
            raw_dis_mesh=stage_data.get("raw_dis_mesh"),
            gt_dis_mesh=stage_data.get("gt_dis_mesh"),
        )
        t_save_pred = time.time()
        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | save_stage_prediction (decode + fix) time: {t_save_pred - t_start:.4f}s")
        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | total stage time: {time.time() - stage_start_time:.4f}s\n")

        if stage_mesh is None:
            break

        current_mesh = stage_mesh
        
    print(f"[{name}] Total process_one_mesh time: {time.time() - t_total_start:.4f}s")


class EMA:
    def __init__(self, model, decay=0.999):
        self.decay = decay
        self.shadow = {}
        self.register(model)

    def register(self, model):
        self.shadow = {}
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()

    def load_state_dict(self, state):
        self.decay = state.get("decay", self.decay)
        self.shadow = {k: v.clone() for k, v in state["shadow"].items()}

    @torch.no_grad()
    def copy_to(self, model):
        """Swap shadow weights into the model in-place (for inference)."""
        copied, missing = 0, 0
        for name, param in model.named_parameters():
            if name in self.shadow:
                param.data.copy_(self.shadow[name].to(device=param.device, dtype=param.dtype))
                copied += 1
            else:
                missing += 1
        print(f"[EMA.copy_to] copied {copied} params into model ({missing} not present in shadow)")


def run_inference_on_all_checkpoints(model, mask_token_id, accelerator, config, device, pad_id, codec, token_shift, ema=None):
    ckpt_dir = "ckpt"
    output_folder = "./pretrain_model_mul_more_data_64_to_256_eval_complex_test_no_fix_holes_no_clean_cfg_1.0_fm"

    # raw_data_txt = "dataset/bpt_res_list.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_retrain_gt/data.txt
    # gt_data_txt = "dataset/trellis_clean_mesh.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_gt/data.txt
    raw_data_txt = "demo/complex_test_data/files/data.txt"
    gt_data_txt = "demo/complex_test_data/files/data.txt"
    # raw_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt"
    # gt_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt"
    stage_resolutions = [64, 128, 256]
    # raw_data_txt = "demo/tripo 1.0/data.txt"
    # gt_data_txt = "demo/tripo 1.0/data.txt"
    
    use_gt_pc=True
    fix_holes=False
    if accelerator.is_main_process:
        print(f"Using raw data txt: {raw_data_txt}")
        print(f"Using gt  data txt: {gt_data_txt}")
        print(f"Searching checkpoints in: {ckpt_dir}")

    # mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=config.data_processing.max_face_num, min_faces=config.data_processing.min_face_num)
    mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=20000, min_faces=500)
    if accelerator.is_main_process:
        print(f"Matched raw/gt pairs: {len(mesh_pairs)}")
        if missing_gt:
            print(f"Raw without matched gt: {len(missing_gt)}")

    search_pattern = os.path.join(ckpt_dir, "checkpoint_*_epoch_*_loss.pth")
    ckpt_files = glob.glob(search_pattern)
    if not ckpt_files:
        if accelerator.is_main_process:
            print("No checkpoints found!")
        return

    ckpt_data = []
    for f in ckpt_files:
        match = re.search(r"checkpoint_(\d+)_epoch_", f)
        if match:
            epoch = int(match.group(1))
            ckpt_data.append((epoch, f))
    ckpt_data.sort(key=lambda x: x[0])

    # stage_resolutions = [32, 64, 128, 256]
    

    # stage_resolutions = [128, 256]

    for epoch, ckpt_path in ckpt_data:
        if accelerator.is_main_process:
            print(f"\n[{'='*40}]")
            print(f"Starting inference for epoch {epoch}")
            print(f"Loading weights from: {ckpt_path}")

        output_folder_ply = os.path.join(output_folder, f"{epoch}_epoch_val_generate_mul")
        if accelerator.is_main_process:
            os.makedirs(output_folder_ply, exist_ok=True)

        checkpoint = torch.load(ckpt_path, map_location=device)
        model_to_load = accelerator.unwrap_model(model)
        # Use strict=False so we can load older ckpts (no null_cond) into the
        # newer model class, and vice versa. Missing/unexpected keys are logged.
        load_result = model_to_load.load_state_dict(checkpoint["model_state_dict"], strict=False)
        missing = list(getattr(load_result, "missing_keys", []) or [])
        unexpected = list(getattr(load_result, "unexpected_keys", []) or [])
        if (missing or unexpected) and accelerator.is_main_process:
            print(f"[load_state_dict] strict=False -> missing={len(missing)} unexpected={len(unexpected)}")
            if missing:
                print(f"  missing: {missing[:8]}{' ...' if len(missing) > 8 else ''}")
            if unexpected:
                print(f"  unexpected: {unexpected[:8]}{' ...' if len(unexpected) > 8 else ''}")

        # Auto-detect EMA: if the ckpt has an `ema_state_dict`, swap those
        # shadow weights into the model for inference. Otherwise fall back to
        # the raw `model_state_dict` we already loaded above. No reliance on
        # config `use_ema` / a separate EMA object — just check the ckpt.
        ema_state = checkpoint.get("ema_state_dict")
        shadow = ema_state.get("shadow") if isinstance(ema_state, dict) else None
        if shadow:
            with torch.no_grad():
                copied, skipped = 0, 0
                for name, param in model_to_load.named_parameters():
                    if name in shadow:
                        param.data.copy_(shadow[name].to(device=param.device, dtype=param.dtype))
                        copied += 1
                    else:
                        skipped += 1
            if accelerator.is_main_process:
                msg = f"[Epoch {epoch}] EMA shadow detected in ckpt -> applied to model ({copied} params"
                if skipped:
                    msg += f", {skipped} model params not in shadow"
                print(msg + ")")
        else:
            if accelerator.is_main_process:
                print(f"[Epoch {epoch}] no EMA in ckpt -> running with raw model_state_dict weights.")

        with torch.no_grad():
            model.eval()
            pbar = tqdm(
                enumerate(mesh_pairs),
                total=len(mesh_pairs),
                desc=f"Infer Epoch {epoch}",
                disable=not accelerator.is_local_main_process,
                dynamic_ncols=True,
            )
            for line_idx, (raw_mesh_path, gt_mesh_path) in pbar:
                if line_idx % accelerator.num_processes != accelerator.process_index:
                    continue
                try:
                    process_one_mesh(
                        raw_mesh_path=raw_mesh_path,
                        gt_mesh_path=gt_mesh_path,
                        model=model,
                        mask_token_id=mask_token_id,
                        codec=codec,
                        device=device,
                        pad_id=pad_id,
                        token_shift=token_shift,
                        config=config,
                        output_folder=output_folder_ply,
                        stage_resolutions=stage_resolutions,
                        fix_holes=fix_holes,
                        use_gt_pc=use_gt_pc,
                    )
                except Exception as e:
                    print(f"[infer error] raw={raw_mesh_path}, gt={gt_mesh_path}: {str(e)}")

        torch.cuda.empty_cache()
        gc.collect()

    if accelerator.is_main_process:
        print("\nBatch inference complete for all checkpoints.")


def main(config):
    print("start init Accelerator...")
    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
    accelerator = Accelerator(
        gradient_accumulation_steps=config.accelerator.gradient_accumulation_steps,
        mixed_precision=config.accelerator.mixed_precision,
        kwargs_handlers=[ddp_kwargs],
    )
    device = accelerator.device

    codec = MeshCodec(args=config.tokenizer)
    codec_config = codec.get_model_config()
    input_num_categories = codec_config["input_num_categories"]
    lm_num_categories = codec_config["output_num_categories"]
    pad_id = codec_config["pad_token_id"]
    token_shift = input_num_categories - 1

    model = DiscreteDiffusionTransformer(
        vocab_A=input_num_categories,
        vocab_B=lm_num_categories,
        max_seq_len=config.data_processing.max_len,
        d_model=config.model.feature_dim,
        encoder_freeze=config.model.encoder_freeze,
        nhead=config.model.num_heads,
        depth=config.model.hourglass_depth,
        dropout=config.model.dropout,
        time_embed_scale=config.training.diffusion_train_steps,
    )
    mask_token_id = model.mask_token_id

    use_ema = getattr(config.training, "use_ema", False)
    ema_decay = getattr(config.training, "ema_decay", 0.999)
    ema = EMA(model, decay=ema_decay) if use_ema else None

    model = accelerator.prepare(model)

    run_inference_on_all_checkpoints(
        model=model,
        mask_token_id=mask_token_id,
        accelerator=accelerator,
        config=config,
        device=device,
        pad_id=pad_id,
        codec=codec,
        token_shift=token_shift,
        ema=ema,
    )


if __name__ == "__main__":
    set_seed(42)
    config = load_config()
    main(config)