import os
import torch
import torch.nn as nn
import argparse
import torch.optim as optim
import logging
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from datetime import datetime
from tqdm import tqdm
import numpy as np

# 假设你的组件路径正确
from models.auto import TemporalTrajectoryPredictor, SpatialGraphRefiner
from models.vae import TripoSGVaeWrapper
from models.clip import FrozenCLIPTextEncoder
from utils.autodataloader import get_mesh_dataloader
from utils.loss import TokenL2Loss

class TextProjectionLayer(nn.Module):
    def __init__(self, in_dim=512, out_dim=64):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(in_dim, in_dim), nn.ReLU(), nn.Linear(in_dim, out_dim))
        self.residual = nn.Linear(in_dim, out_dim) if in_dim != out_dim else nn.Identity()
    def forward(self, x): return self.net(x) + self.residual(x)

def setup_logger(save_dir, rank):
    os.makedirs(save_dir, exist_ok=True)
    logger = logging.getLogger("TrainLogger")
    if rank == 0: # 仅主进程记录日志到文件
        logger.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s - %(message)s')
        file_handler = logging.FileHandler(os.path.join(save_dir, f"train_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"))
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
    else:
        logger.setLevel(logging.WARNING) # 子进程仅记录警告
    return logger

def train(args):
    # 1. 初始化分布式环境
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()

    logger = setup_logger(args.save_dir, rank)
    
    # 2. 初始化模型
    vae = TripoSGVaeWrapper(device=device)
    vae.eval()
    for param in vae.parameters(): param.requires_grad = False
    
    text_encoder = FrozenCLIPTextEncoder(args.clip_model_path, device=device)
    text_proj = TextProjectionLayer(in_dim=512, out_dim=args.token_dim).to(device)
    temporal_model = TemporalTrajectoryPredictor(num_tokens=args.num_tokens, token_dim=args.token_dim, max_seq_len=16).to(device)
    spatial_model = SpatialGraphRefiner(num_tokens=args.num_tokens, token_dim=args.token_dim, text_dim=args.token_dim).to(device)
    
    # 包装 DDP
    text_proj = DDP(text_proj, device_ids=[local_rank])
    temporal_model = DDP(temporal_model, device_ids=[local_rank])
    spatial_model = DDP(spatial_model, device_ids=[local_rank])
    
    optimizer = optim.Adam(
        list(temporal_model.parameters()) + list(spatial_model.parameters()) + list(text_proj.parameters()), 
        lr=args.lr
    )
    
    criterion = TokenL2Loss()
    dataloader = get_mesh_dataloader(batch_size=args.batch_size, num_workers=4)
    
    # 3. 训练循环
    for epoch in range(args.epochs):
        temporal_model.train()
        spatial_model.train()
        text_proj.train()
        
        # 进度条仅主进程显示
        pbar = tqdm(dataloader, desc=f"Epoch {epoch+1:03d}", disable=(rank != 0), ncols=150)
        
        for batch in pbar:
            optimizer.zero_grad()
            vertices, faces = batch['vertices'], batch['faces']
            
            with torch.no_grad():
                latents_batch = []
                for i in range(len(vertices)):
                    sample_latents = []
                    for t in range(vertices[i].shape[0]):
                        v_t = torch.tensor(vertices[i][t], dtype=torch.float32, device=device)
                        f_t = torch.tensor(faces[i], dtype=torch.long, device=device)
                        l = vae.encode_mesh(v_t, f_t)
                        if isinstance(l, np.ndarray): l = torch.from_numpy(l).to(device)
                        sample_latents.append(l.squeeze(0))
                    latents_batch.append(torch.stack(sample_latents, dim=0))
                latents = torch.stack(latents_batch, dim=0)
            
            text_embed = text_proj(text_encoder(list(batch['caption']))).unsqueeze(1)
            t_idx = torch.randint(1, 16, (1,)).item()
            
            history = latents[:, :t_idx, :, :]
            target = latents[:, t_idx, :, :]
            
            next_t_init = temporal_model(history)
            next_t_refined = spatial_model(next_t_init, text_embed, torch.zeros((args.batch_size, args.num_tokens, args.num_tokens)).to(device))
            
            total_loss = criterion(next_t_init, target) + criterion(next_t_refined, target)
            total_loss.backward()
            optimizer.step()
            
            if rank == 0:
                pbar.set_postfix({"Loss": f"{total_loss.item():.3f}"})
        
        # 4. 主进程保存模型
        if rank == 0 and (epoch + 1) % args.save_every_k_epochs == 0:
            ckpt_path = os.path.join(args.save_dir, f"model_epoch_{epoch+1}.pth")
            torch.save({
                'temporal': temporal_model.module.state_dict(),
                'spatial': spatial_model.module.state_dict(),
                'proj': text_proj.module.state_dict()
            }, ckpt_path)
            logger.info(f"模型已保存: {ckpt_path}")

    dist.destroy_process_group()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # ... 原有参数 ...
    parser.add_argument("--num_tokens", type=int, default=2048)
    parser.add_argument("--token_dim", type=int, default=64)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--save_every_k_epochs", type=int, default=10)
    parser.add_argument("--save_dir", type=str, default="./checkpoints")
    parser.add_argument("--clip_model_path", type=str, default="/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/pretrain/clip-vit-base-patch32")
    args = parser.parse_args()
    train(args)