"""
Generate predicted latent z for offline evaluation.

This script uses an MDM model trained on preprocessed_posterior data to generate
predicted z values and saves them as .pt files for later decoding and evaluation
in the SAME project.

It loads text and length information from preprocessed posterior files, using only
the 'text' and 'm_len' fields.

Workflow:
1. Load the MDM model
2. Iterate over all test samples and generate predicted z from text (512-dim latent)
3. Save each sample's z_pred, text information, and length

Output format:
    {
        'z_pred': torch.Tensor,  # shape (m_len, 512), predicted feature z
        'text': {
            'caption': str,      # text description
            'tokens': list,      # token list
            'start_time': float,
            'end_time': float
        },
        'm_len': int             # motion length
    }

Output directory structure:
    output_dir/
        repeat_1/
            posterior_000000.pt
            posterior_000001.pt
            ...
        repeat_2/
            ...
        config.json              # saved generation config

Usage example:
    python -m eval.eval_humanml_preproc \
        --model_path ./save/first_exp_train/model000100000.pt \
        --posterior_dir /path/to/val_posteriors \
        --output_dir ./output/eval_first_exp \
        --num_repetitions 20 \
        --guidance_param 2.5
"""

import os
import sys
import json
import torch
import numpy as np
from datetime import datetime
from argparse import ArgumentParser
from tqdm import tqdm
from pathlib import Path
from torch.utils.data import Dataset, DataLoader

# MDM imports
from utils.fixseed import fixseed
from utils.model_util import create_model_and_diffusion, load_saved_model
from utils import dist_util
from diffusion import logger
from utils.sampler_util import ClassifierFreeSampleModel

torch.multiprocessing.set_sharing_strategy('file_system')


def evaluation_preproc_args():
    """Parse evaluation arguments."""
    parser = ArgumentParser(description='Generate predicted z for offline evaluation')
    
    # MDM model arguments
    parser.add_argument("--model_path", required=True, type=str,
                        help="MDM model checkpoint path")
    parser.add_argument("--guidance_param", default=2.5, type=float,
                        help="Classifier-free guidance scale")
    parser.add_argument("--use_ema", action='store_true',
                        help="Use the EMA model")
    
    # Data source arguments
    parser.add_argument("--posterior_dir", required=True, type=str,
                        help="Directory containing preprocessed posterior files")
    
    # Output arguments
    parser.add_argument("--output_dir", required=True, type=str,
                        help="Output directory")
    
    # Evaluation arguments
    parser.add_argument("--num_repetitions", default=20, type=int,
                        help="Number of repetitions")
    parser.add_argument("--seed", default=10, type=int, help="Random seed")
    parser.add_argument("--device", default=0, type=int, help="GPU device ID")
    parser.add_argument("--batch_size", default=32, type=int, 
                        help="Batch size")
    
    args = parser.parse_args()
    return args


def load_mdm_args(model_path):
    """Load training arguments from the checkpoint directory."""
    args_path = os.path.join(os.path.dirname(model_path), 'args.json')
    with open(args_path, 'r') as f:
        args_dict = json.load(f)
    
    # Create a simple object to store arguments.
    class Args:
        pass
    
    args = Args()
    for k, v in args_dict.items():
        setattr(args, k, v)
    
    return args


class PosteriorTextDataset(Dataset):
    """
    Dataset that loads text and length information from preprocessed posterior files.
    Uses only the 'text' and 'm_len' fields and ignores 'mean', 'logvar', and 'sample'.
    """
    def __init__(self, posterior_dir):
        """
        Args:
            posterior_dir: directory containing preprocessed posterior files
        """
        self.posterior_dir = posterior_dir
        
        # Scan the directory and collect all posterior file paths.
        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']
        ])
        
        self.file_paths = [os.path.join(posterior_dir, f) for f in posterior_files]
        
        print(f"[PosteriorTextDataset] Loaded {len(self.file_paths)} samples from {posterior_dir}")
    
    def __len__(self):
        return len(self.file_paths)
    
    def __getitem__(self, idx):
        """
        Load text and length information for sample idx.
        
        Returns:
            dict: contains fields such as inp, text, and lengths
        """
        posterior_data = torch.load(self.file_paths[idx])
        
        # Get text information.
        text_info = posterior_data['text']
        m_len = posterior_data['m_len']
        
        # Handle text information across multiple formats.
        if isinstance(text_info, dict):
            caption = text_info.get('caption', '')
            tokens = text_info.get('tokens', [])
        else:
            caption = str(text_info)
            tokens = []
        
        # Convert the token list to a string joined with '_'.
        tokens_str = '_'.join(tokens) if tokens else ''
        
        # Create a dummy inp; only shape information is needed.
        # Shape: (C, 1, T), where C=512 and T=m_len.
        inp = torch.zeros(512, 1, m_len)
        
        # Extract the filename without path.
        filename = os.path.basename(self.file_paths[idx])
        
        return {
            'inp': inp,
            'text': caption,
            'tokens': tokens_str,
            'lengths': m_len,
            'text_info': text_info,  # Save full text_info for output.
            'filename': filename,  # Add filename.
        }


def collate_posterior_text(batch):
    """
    Collate function for PosteriorTextDataset
    Output format is compatible with MDM collate.
    """
    # Find the maximum length.
    max_len = max(item['lengths'] for item in batch)
    
    batch_size = len(batch)
    
    # Create padded tensors.
    # inp: (B, C, 1, T_max)
    inp_padded = torch.zeros(batch_size, 512, 1, max_len)
    
    # mask: (B, 1, 1, T_max)
    mask = torch.zeros(batch_size, 1, 1, max_len).bool()
    
    lengths = []
    texts = []
    tokens_list = []
    text_infos = []
    filenames = []
    
    for i, item in enumerate(batch):
        m_len = item['lengths']
        lengths.append(m_len)
        texts.append(item['text'])
        tokens_list.append(item['tokens'])
        text_infos.append(item['text_info'])
        filenames.append(item['filename'])
        
        # Set mask; valid positions are True.
        mask[i, :, :, :m_len] = True
    
    lengths = torch.tensor(lengths)
    
    # Build model_kwargs.
    cond = {
        'y': {
            'mask': mask,
            'lengths': lengths,
            'text': texts,
            'tokens': tokens_list,
            'text_infos': text_infos,  # Also save full text_info.
            'filenames': filenames,  # Add filename list.
        }
    }
    
    return inp_padded, cond


def get_posterior_text_loader(posterior_dir, batch_size=32, shuffle=False, 
                               num_workers=4):
    """
    Create a DataLoader that loads text information from posterior files.
    """
    dataset = PosteriorTextDataset(posterior_dir)
    
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=collate_posterior_text,
        pin_memory=True,
        drop_last=False,
    )
    
    return dataloader


class MDMLatentGenerator:
    """
    Latent generator for the MDM model.
    Wraps the MDM model and provides an interface for generating z.
    """
    def __init__(self, model, diffusion, args, scale=2.5):
        self.model = model
        self.diffusion = diffusion
        self.args = args
        self.scale = scale
        self.sample_fn = diffusion.p_sample_loop
        
    def generate(self, model_kwargs, motion_shape):
        """
        Generate latent z.
        
        Args:
            model_kwargs: dictionary containing conditioning information
            motion_shape: output shape (B, njoints, nfeats, T)
            
        Returns:
            sample: generated latent z with shape (B, T, C)
        """
        # Add CFG scale.
        if self.scale != 1.:
            model_kwargs['y']['scale'] = torch.ones(motion_shape[0], 
                                                     device=dist_util.dev()) * self.scale
        
        # Encode text once.
        if 'text' in model_kwargs['y'].keys() and 'text_embed' not in model_kwargs['y'].keys():
            model_kwargs['y']['text_embed'] = self.model.encode_text(model_kwargs['y']['text'])
        
        # Generate.
        sample = self.sample_fn(
            self.model,
            motion_shape,
            clip_denoised=False,
            model_kwargs=model_kwargs,
            skip_timesteps=0,
            init_image=None,
            progress=False,
            dump_steps=None,
            noise=None,
            const_noise=False,
        )
        
        # Convert shape: (B, C, 1, T) -> (B, T, C).
        sample = sample.squeeze(2).permute(0, 2, 1)  # (B, T, C)
        
        return sample


def generate_and_save_latents(mdm_generator, dataloader, output_dir, repeat_idx,
                               progress=True):
    """
    Generate latent z with MDM and save it.
    
    Args:
        mdm_generator: MDM model generator
        dataloader: data loader that provides text conditions
        output_dir: output directory
        repeat_idx: current repetition index, starting from 1
        progress: whether to show a progress bar
        
    Returns:
        num_generated: number of generated samples
    """
    # Create the repetition directory.
    repeat_dir = os.path.join(output_dir, f'repeat_{repeat_idx}')
    os.makedirs(repeat_dir, exist_ok=True)
    
    real_num_batches = len(dataloader)
    
    desc = f'Generating repeat_{repeat_idx}'
    iterator = tqdm(enumerate(dataloader), total=real_num_batches, desc=desc) if progress else enumerate(dataloader)
    
    sample_idx = 0
    
    with torch.no_grad():
        for i, (motion, model_kwargs) in iterator:
            # Move to device.
            model_kwargs['y'] = {
                key: val.to(dist_util.dev()) if torch.is_tensor(val) else val 
                for key, val in model_kwargs['y'].items()
            }
            
            # Get tokens.
            tokens = [t.split('_') for t in model_kwargs['y']['tokens']]
            
            # Determine generation shape.
            batch_size = motion.shape[0]
            n_frames = motion.shape[-1]
            motion_shape = (batch_size, mdm_generator.model.njoints, mdm_generator.model.nfeats, n_frames)
            
            # Generate latent z.
            latent_z = mdm_generator.generate(model_kwargs, motion_shape)  # (B, T, C)
            
            # Get valid frame counts.
            lengths = model_kwargs['y']['lengths'].cpu().numpy()
            
            # Save each sample.
            for bs_i in range(batch_size):
                m_len = int(lengths[bs_i])
                
                # Slice the valid part of z_pred.
                z_pred = latent_z[bs_i, :m_len, :].cpu()  # (m_len, 512)
                
                # Use the full text_info from the posterior file.
                text_info = model_kwargs['y']['text_infos'][bs_i]
                
                # Get the original filename.
                original_filename = model_kwargs['y']['filenames'][bs_i]
                
                # Build the save dictionary.
                save_dict = {
                    'z_pred': z_pred,
                    'text': text_info,
                    'm_len': m_len
                }
                
                # Save the file using the original filename.
                save_path = os.path.join(repeat_dir, original_filename)
                torch.save(save_dict, save_path)
                
                sample_idx += 1
    
    return sample_idx


def main():
    args = evaluation_preproc_args()
    fixseed(args.seed)
    
    replication_times = args.num_repetitions
    print(f'Number of repetitions: {replication_times}')
    
    # Create the output directory.
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Set up distributed environment.
    dist_util.setup_dist(args.device)
    logger.configure()
    
    # =========================================
    # 1. Load the MDM model
    # =========================================
    logger.log("Loading MDM model arguments...")
    mdm_args = load_mdm_args(args.model_path)
    
    # Ensure required arguments exist.
    if not hasattr(mdm_args, 'pred_len'):
        mdm_args.pred_len = 0
        mdm_args.context_len = 0
    if not hasattr(mdm_args, 'autoregressive'):
        mdm_args.autoregressive = False
    
    logger.log("Creating data loader...")
    
    # Use preprocessed posterior files, using only text and m_len.
    logger.log(f"Using posterior files from: {args.posterior_dir}")
    gen_loader = get_posterior_text_loader(
        posterior_dir=args.posterior_dir,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=4
    )
    
    logger.log("Creating MDM model and diffusion...")
    # Create the model directly; no dummy loader is needed.
    model, diffusion = create_model_and_diffusion(mdm_args, None)
    
    logger.log(f"Loading MDM checkpoints from [{args.model_path}]...")
    load_saved_model(model, args.model_path, use_avg=args.use_ema)
    
    if args.guidance_param != 1:
        model = ClassifierFreeSampleModel(model)
    model.to(dist_util.dev())
    model.eval()
    
    # =========================================
    # 2. Create generator
    # =========================================
    mdm_generator = MDMLatentGenerator(model, diffusion, mdm_args, scale=args.guidance_param)
    
    # =========================================
    # 3. Save configuration information
    # =========================================
    num_samples = len(gen_loader.dataset)
    config = {
        'model_path': args.model_path,
        'guidance_param': args.guidance_param,
        'use_ema': args.use_ema,
        'num_samples': num_samples,
        'num_repetitions': replication_times,
        'seed': args.seed,
        'batch_size': args.batch_size,
        'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        'njoints': mdm_args.njoints,
        'latent_dim': mdm_args.latent_dim,
        'posterior_dir': args.posterior_dir,
    }
    
    config_path = os.path.join(args.output_dir, 'config.json')
    with open(config_path, 'w') as f:
        json.dump(config, f, indent=4)
    print(f'Config saved to {config_path}')
    
    # =========================================
    # 4. Generate and save latent z
    # =========================================
    logger.log("Starting generation...")
    
    for repeat_idx in range(1, replication_times + 1):
        print(f'\n{"="*60}')
        print(f'Repetition {repeat_idx}/{replication_times}')
        print(f'{"="*60}')
        
        # Set a different random seed for each repetition.
        fixseed(args.seed + repeat_idx - 1)
        
        num_generated = generate_and_save_latents(
            mdm_generator=mdm_generator,
            dataloader=gen_loader,
            output_dir=args.output_dir,
            repeat_idx=repeat_idx,
            progress=True
        )
        
        print(f'Generated {num_generated} samples for repeat_{repeat_idx}')
    
    print(f'\n{"="*60}')
    print(f'Generation completed!')
    print(f'Output directory: {args.output_dir}')
    print(f'Total repetitions: {replication_times}')
    print(f'Samples per repetition: {num_generated}')
    print(f'{"="*60}')


if __name__ == '__main__':
    main()