import os
import torch
import torch.nn as nn
import argparse
import torch.optim as optim
import logging
from datetime import datetime
from tqdm import tqdm

# 导入您定义的组件
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):
    """两层线性层 + 残差连接：将 512 维投影至 64 维"""
    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):
    os.makedirs(save_dir, exist_ok=True)
    time_str = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = os.path.join(save_dir, f"train_{time_str}.log")
    
    logger = logging.getLogger("TrainLogger")
    logger.setLevel(logging.INFO)
    
    formatter = logging.Formatter('%(asctime)s - %(message)s')
    
    file_handler = logging.FileHandler(log_file)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    
    return logger

def count_parameters(model, name, logger):
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    logger.info(f"[{name}] 可训练参数量: {trainable:,}")
    return trainable

def train(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger = setup_logger(args.save_dir)
    
    # 1. 初始化模型
    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)
    
    logger.info("--- 参数统计 ---")
    count_parameters(temporal_model, "TemporalModel", logger)
    count_parameters(spatial_model, "SpatialModel", logger)
    count_parameters(text_proj, "TextProjection", logger)
    
    criterion = TokenL2Loss()
    optimizer = optim.Adam(
        list(temporal_model.parameters()) + list(spatial_model.parameters()) + list(text_proj.parameters()), 
        lr=args.lr
    )
    
    dataloader = get_mesh_dataloader(batch_size=args.batch_size, num_workers=4)
    
    # 定义美观的进度条格式
    bar_format = '{l_bar}{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
    
    # 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}/{args.epochs:03d}", 
                    ncols=300, bar_format=bar_format)
        
        for batch in pbar:
            optimizer.zero_grad()
            
            vertices = batch['vertices'].to(device)
            faces = batch['faces'].to(device)

            max_v_idx = vertices.shape[2] - 1
            if faces.max() > max_v_idx:
                pbar.write(f"\n[!] 发现非法数据，索引 {faces.max()} 超出范围，跳过Batch")
                continue
            
            with torch.no_grad():
                B, T, V, C = vertices.shape
                latents = []
                for t in range(T):
                    l = vae.encode_mesh(vertices[:, t].squeeze(0), faces.squeeze(0))
                    latents.append(l)
                latents = torch.stack(latents, dim=1)
            
            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))
            
            loss_t = criterion(next_t_init, target)
            loss_s = criterion(next_t_refined, target)
            total_loss = loss_t + loss_s
            
            total_loss.backward()
            optimizer.step()
            
            pbar.set_postfix({
                "Loss": f"{total_loss.item():.3f}",
                "Temp": f"{loss_t.item():.3f}",
                "Spat": f"{loss_s.item():.3f}"
            })
        
        logger.info(f"Epoch {epoch+1}/{args.epochs} 完成 - Total Loss: {total_loss.item():.4f}")
        
        if (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.state_dict(),
                'spatial': spatial_model.state_dict(),
                'proj': text_proj.state_dict()
            }, ckpt_path)
            logger.info(f"模型已保存至: {ckpt_path}")

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)