src/mdm/data_loaders/preprocessed_posterior_loader.py

"""
DataLoader for loading data from preprocessed posterior files.
Adapted to this repository's training flow and compatible with train_mdm.py.
Based on the external preprocessed_posterior_loader.py.

The output format matches the original t2m_collate and can be used directly by existing training code.
"""

import os
import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
from pathlib import Path
from os.path import join as pjoin
from data_loaders.tensors import collate as mdm_collate


class PreprocessedPosteriorDataset(Dataset):
    """
    Dataset for loading data from preprocessed posterior files.
    Adapts the data format to the MDM training flow.
    
    The output format matches datasets such as HumanML3D and can use the existing collate function directly.
    """
    def __init__(self, posterior_dir, max_samples=None, resample=False, 
                 mean=None, std=None):
        """
        Args:
            posterior_dir: directory containing preprocessed posterior files
            max_samples: maximum number of samples; None loads all samples
            resample: whether to resample the latent while loading; False uses the preprocessed sample
            mean: mean used for normalization (optional, for compatibility)
            std: standard deviation used for normalization (optional, for compatibility)
        """
        self.posterior_dir = posterior_dir
        self.resample = resample
        
        # Scan the directory and collect posterior file paths, excluding meta_info.pt and preprocess_config.pt.
        all_files = os.listdir(posterior_dir)
        posterior_files = sorted([
            f for f in all_files 
            if f.endswith('.pt') and f not in ['meta_info.pt', 'preprocess_config.pt']
        ])
        
        # Save the file path list.
        self.file_paths = [pjoin(posterior_dir, f) for f in posterior_files]
        self.total_samples = len(self.file_paths)
        
        print(f"Scanned directory: found {self.total_samples} posterior files")
        print(f"Scanned directory: found {self.total_samples} posterior files")
        
        # Limit the number of samples.
        if max_samples is not None:
            self.total_samples = min(self.total_samples, max_samples)
            self.file_paths = self.file_paths[:max_samples]
            print(f"Limited to the first {self.total_samples} samples")
        
        # Try to load metadata if available.
        meta_path = pjoin(posterior_dir, 'meta_info.pt')
        if os.path.exists(meta_path):
            meta_info = torch.load(meta_path)
            self.config = meta_info.get('config', {})
            print("Loaded metadata configuration")
        else:
            self.config = {}
        
        # Store mean and std for compatibility with dataset interfaces such as HumanML3D.
        self.mean = mean
        self.std = std
        
        print(f"PreprocessedPosteriorDataset initialized with {self.total_samples} samples")
        if self.resample:
            print("Resampling mode enabled: resampling latent from posterior on each load")
    
    def sample_from_posterior(self, mean, logvar):
        """
        Sample from the posterior distribution.
        Uses the reparameterization trick: z = mean + std * eps.
        
        Args:
            mean: posterior mean
            logvar: posterior log variance
        
        Returns:
            sampled latent
        """
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mean + eps * std
    
    def inv_transform(self, data):
        """Inverse transform for compatibility."""
        if self.mean is not None and self.std is not None:
            return data * self.std + self.mean
        return data
    
    def __len__(self):
        return self.total_samples
    
    def __getitem__(self, idx):
        """
        Load posterior data for sample idx.
        
        Returns:
            dict: matches the format expected by the original t2m_collate and contains:
                - inp: motion data with shape (C, 1, T)
                - text: text description (string)
                - tokens: token list if available
                - lengths: motion length (integer)
        """
        posterior_path = self.file_paths[idx]
        
        if not os.path.exists(posterior_path):
            raise FileNotFoundError(f"Posterior file not found: {posterior_path}")
        
        posterior_data = torch.load(posterior_path)
        
        # Get data.
        text_info = posterior_data['text']
        m_len = posterior_data['m_len']
        
        # If resampling is enabled, resample from mean and logvar.
        if self.resample:
            mean = posterior_data['mean']      # (T, C)
            logvar = posterior_data['logvar']  # (T, C)
            sample = self.sample_from_posterior(mean, logvar)
        else:
            if 'sample' in posterior_data:
                sample = posterior_data['sample']  # (T, C)
            elif 'z' in posterior_data:
                sample = posterior_data['z']       # (T, C)
            else:
                raise KeyError(f"Posterior file is missing the 'sample' or 'z' key: {posterior_path}")
        
        # Convert to the format expected by MDM: (T, C) -> (C, 1, T).
        # Align with the conversion in t2m_collate: torch.tensor(b[4].T).float().unsqueeze(1).
        inp = sample.T.unsqueeze(1).float()  # (C, 1, T)
        
        # Handle length.
        if isinstance(m_len, (int, np.integer)):
            length = int(m_len)
        elif hasattr(m_len, 'item'):
            length = int(m_len.item())
        else:
            length = int(m_len)
        
        # Handle text information across multiple formats.
        # text_info may be a string, dict, or another format.
        if isinstance(text_info, dict):
            text = text_info.get('caption', str(text_info))
            tokens = text_info.get('tokens', None)
        elif isinstance(text_info, str):
            text = text_info
            tokens = None
        else:
            text = str(text_info)
            tokens = None
        
        # Build the return dictionary to match the format expected by t2m_collate.
        result = {
            'inp': inp,
            'text': text,
            'lengths': length,
        }
        
        # Add tokens if available.
        if tokens is not None:
            result['tokens'] = tokens
        
        return result


def collate_preprocessed_posterior(batch):
    """
    Use the original collate function directly to keep the output format identical.
    
    Args:
        batch: list of dicts, each containing inp, text, lengths, etc.
    
    Returns:
        motion: (B, C, 1, T_max) tensor
        cond: dict containing a 'y' key with mask, lengths, text, etc.
        
    Output format exactly matches t2m_collate:
        - motion: (B, C, 1, T_max)
        - cond['y']['mask']: (B, 1, 1, T_max)
        - cond['y']['lengths']: (B,)
        - cond['y']['text']: list of str
    """
    # Use the original collate function directly.
    return mdm_collate(batch)


class PreprocessedPosteriorDataLoader:
    """
    DataLoader wrapper for preprocessed posterior data.
    Provides an interface compatible with existing data loaders.
    
    This class emulates the original DataLoader interface, including dataset attribute access.
    """
    def __init__(self, posterior_dir, batch_size=32, shuffle=True, 
                 num_workers=4, max_samples=None, resample=False,
                 mean=None, std=None, drop_last=True):
        """
        Create a DataLoader for preprocessed posterior data.
        
        Args:
            posterior_dir: directory containing preprocessed posterior files
            batch_size: batch size
            shuffle: whether to shuffle
            num_workers: number of data loading workers
            max_samples: maximum number of samples
            resample: whether to resample the latent while loading
            mean: mean used for normalization
            std: standard deviation used for normalization
            drop_last: whether to drop the last incomplete batch
        """
        self.dataset = PreprocessedPosteriorDataset(
            posterior_dir, 
            max_samples=max_samples, 
            resample=resample,
            mean=mean,
            std=std
        )
        
        self.dataloader = DataLoader(
            self.dataset,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=num_workers,
            collate_fn=collate_preprocessed_posterior,
            pin_memory=True,
            drop_last=drop_last,
        )
        
        # Store attributes for external access and training_loop.py compatibility.
        self.batch_size = batch_size
        # These attributes are used for compatibility with target_cond_modifier in training_loop.
        # preprocessed_posterior may not need these, but keep the interface.
        self.mean = mean
        self.std = std
    
    def __iter__(self):
        return iter(self.dataloader)
    
    def __len__(self):
        return len(self.dataloader)


def get_preprocessed_posterior_loader(posterior_dir, batch_size=32, shuffle=True,
                                       num_workers=4, max_samples=None, resample=False,
                                       mean=None, std=None, drop_last=True):
    """
    Convenience function for creating a DataLoader for preprocessed posterior data.
    
    Args:
        posterior_dir: directory containing preprocessed posterior files
        batch_size: batch size
        shuffle: whether to shuffle
        num_workers: number of data loading workers
        max_samples: maximum number of samples
        resample: whether to resample the latent while loading
        mean: mean used for normalization
        std: standard deviation used for normalization
        drop_last: whether to drop the last incomplete batch
    
    Returns:
        PreprocessedPosteriorDataLoader object
    """
    return PreprocessedPosteriorDataLoader(
        posterior_dir=posterior_dir,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        max_samples=max_samples,
        resample=resample,
        mean=mean,
        std=std,
        drop_last=drop_last
    )
    


# Test code
if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='Test the preprocessed posterior data loader (MDM-compatible version)')
    parser.add_argument('--posterior_dir', type=str, required=True,
                        help='Directory containing preprocessed posterior files')
    parser.add_argument('--batch_size', type=int, default=4,
                        help='Batch size')
    parser.add_argument('--max_samples', type=int, default=None,
                        help='Maximum number of samples for testing')
    parser.add_argument('--resample', action='store_true',
                        help='Enable resampling mode')
    args = parser.parse_args()
    
    print("="*50)
    print("Testing preprocessed posterior data loader (MDM-compatible version)")
    print("="*50)
    
    # Create dataloader.
    dataloader = get_preprocessed_posterior_loader(
        args.posterior_dir,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=0,  # Use a single process for testing.
        max_samples=args.max_samples,
        resample=args.resample
    )
    
    print(f"\nDataLoader created with {len(dataloader)} batches")
    print(f"Resampling mode: {'enabled' if args.resample else 'disabled'}")
    
    # Test loading the first batch.
    print("\nTesting first batch load...")
    for batch_idx, (motion, cond) in enumerate(dataloader):
        print(f"\nBatch {batch_idx}:")
        print(f"  motion shape: {motion.shape}")  # Expected: (B, C, 1, T_max)
        print(f"  mask shape: {cond['y']['mask'].shape}")  # Expected: (B, 1, 1, T_max)
        print(f"  lengths: {cond['y']['lengths']}")
        print(f"  texts: {len(cond['y']['text'])} text entries")
        
        # Only test the first batch.
        if batch_idx == 0:
            print("\nText information for the first sample:")
            print(f"  {cond['y']['text'][0]}")
            
            # Show motion statistics.
            print("\nMotion statistics for the first sample:")
            first_len = cond['y']['lengths'][0].item()
            first_motion_valid = motion[0, :, :, :first_len]  # Only inspect the valid part.
            print(f"  valid length: {first_len}")
            print(f"  mean: {first_motion_valid.mean().item():.6f}")
            print(f"  std: {first_motion_valid.std().item():.6f}")
            print(f"  min: {first_motion_valid.min().item():.6f}")
            print(f"  max: {first_motion_valid.max().item():.6f}")
            
            break
    
    print("\nTest completed!")
    print("\nUsage example:")
    print("  from data_loaders.preprocessed_posterior_loader import get_preprocessed_posterior_loader")
    print("  data = get_preprocessed_posterior_loader(posterior_dir, batch_size=32)")
    print("  for motion, cond in data:")
    print("      # motion: (B, C, 1, T)")
    print("      # cond['y']['mask']: (B, 1, 1, T)")
    print("      # cond['y']['lengths']: (B,)")
    print("      # cond['y']['text']: list of texts")