"""Pipeline orchestrator for multi-step training workflows.

Manages sequential execution of training pipeline steps with model freezing,
loss-threshold transitions, and checkpoint management.
"""

from dataclasses import asdict
from typing import Any, Optional

import torch.nn as nn
from fluxflow.utils import get_logger, safe_vae_sample, save_sample_images

from .checkpoint_manager import CheckpointManager
from .pipeline_config import PipelineConfig, PipelineStepConfig

logger = get_logger(__name__)


class TrainingPipelineOrchestrator:
    """
    Orchestrates multi-step training pipelines.

    Manages:
    - Sequential step execution
    - Model freeze/unfreeze per step
    - Loss-threshold monitoring for transitions
    - Checkpoint save/load with pipeline metadata
    - Resume from any step

    Example:
        >>> config = parse_pipeline_config(config_dict)
        >>> orchestrator = TrainingPipelineOrchestrator(
        ...     config=config,
        ...     models=models_dict,
        ...     checkpoint_manager=checkpoint_mgr,
        ...     accelerator=accelerator,
        ... )
        >>> orchestrator.run()
    """

    def __init__(
        self,
        pipeline_config: PipelineConfig = None,
        checkpoint_manager: CheckpointManager = None,
        accelerator: Any = None,
        device: Any = None,
        # Legacy signature support (for tests)
        config: PipelineConfig = None,
        models: dict[str, nn.Module] = None,
        dataloader: Any = None,
        dataset: Any = None,
    ):
        """
        Initialize pipeline orchestrator.

        Args:
            pipeline_config: Parsed pipeline configuration (new signature)
            checkpoint_manager: Checkpoint manager instance (new signature)
            accelerator: Accelerate accelerator instance (new signature)
            device: Target device (new signature)
            config: Parsed pipeline configuration (legacy, for tests)
            models: Dictionary of model components (legacy, for tests)
            dataloader: Training dataloader (legacy, for tests)
            dataset: Training dataset (legacy, for tests)
        """
        # Support both new and legacy signatures
        self.config = pipeline_config or config
        self.checkpoint_manager = checkpoint_manager
        self.accelerator = accelerator
        self.device = device

        # Legacy support
        self.models = models or {}
        self.dataloader = dataloader
        self.dataset = dataset

        # Pipeline state
        self.current_step_index = 0
        self.global_step = 0
        self.steps_completed: list[str] = []

        # Metric tracking for loss-threshold transitions
        self.step_metrics: dict[str, dict[str, list[float]]] = {}

        # Validate models dictionary if provided (legacy mode)
        if self.models:
            self._validate_models()

    def _validate_models(self) -> None:
        """Validate that all required model components are present."""
        required = {"compressor", "expander", "flow_processor", "text_encoder", "discriminator"}
        missing = required - set(self.models.keys())
        if missing:
            logger.warning(f"Missing model components (may be provided to run()): {missing}")

    def freeze_model(self, model_name: str) -> None:
        """
        Freeze model parameters.

        Args:
            model_name: Name of model to freeze (e.g., 'compressor', 'text_encoder')
        """
        if model_name not in self.models:
            logger.warning(f"Cannot freeze '{model_name}': not found in models dict")
            return

        model = self.models[model_name]
        for param in model.parameters():
            param.requires_grad = False
        model.eval()

        # Count frozen parameters
        num_params = sum(p.numel() for p in model.parameters())
        logger.info(f"✓ Frozen: {model_name} ({num_params:,} parameters)")

    def unfreeze_model(self, model_name: str) -> None:
        """
        Unfreeze model parameters.

        Args:
            model_name: Name of model to unfreeze
        """
        if model_name not in self.models:
            logger.warning(f"Cannot unfreeze '{model_name}': not found in models dict")
            return

        model = self.models[model_name]
        for param in model.parameters():
            param.requires_grad = True
        model.train()

        # Count unfrozen parameters
        num_params = sum(p.numel() for p in model.parameters())
        logger.info(f"✓ Unfrozen: {model_name} ({num_params:,} parameters)")

    def configure_step_models(
        self, step: PipelineStepConfig, models: dict[str, nn.Module] = None
    ) -> None:
        """
        Configure models for pipeline step (freeze/unfreeze).

        Args:
            step: Pipeline step configuration
            models: Dictionary of models (optional, uses self.models if not provided)
        """
        models_dict = models or self.models
        logger.info(f"Configuring models for step '{step.name}'...")

        # Freeze specified models
        for model_name in step.freeze:
            if model_name not in models_dict:
                logger.warning(f"Cannot freeze '{model_name}': not found in models dict")
                continue
            model = models_dict[model_name]
            for param in model.parameters():
                param.requires_grad = False
            logger.info(f"Frozen model: {model_name}")

        # Unfreeze specified models
        for model_name in step.unfreeze:
            if model_name not in models_dict:
                logger.warning(f"Cannot unfreeze '{model_name}': not found in models dict")
                continue
            model = models_dict[model_name]
            for param in model.parameters():
                param.requires_grad = True
            logger.info(f"Unfrozen model: {model_name}")

        # Log final state
        if models_dict:
            trainable_params = sum(
                p.numel() for m in models_dict.values() for p in m.parameters() if p.requires_grad
            )
            total_params = sum(p.numel() for m in models_dict.values() for p in m.parameters())
            frozen_params = total_params - trainable_params

            if total_params > 0:
                logger.info(
                    f"Model configuration complete: "
                    f"{trainable_params:,} trainable, {frozen_params:,} frozen "
                    f"({100.0 * trainable_params / total_params:.1f}% trainable)"
                )
            else:
                logger.warning("No model parameters found")

    def update_metrics(self, step_name: str, losses: dict[str, float]) -> None:
        """
        Update metric history for transition monitoring.

        Args:
            step_name: Name of current step
            losses: Dictionary of loss values from training step
        """
        if step_name not in self.step_metrics:
            self.step_metrics[step_name] = {}

        for metric_name, value in losses.items():
            if metric_name not in self.step_metrics[step_name]:
                self.step_metrics[step_name][metric_name] = []

            # Keep last 100 values for smoothing
            history = self.step_metrics[step_name][metric_name]
            history.append(float(value))
            if len(history) > 100:
                history.pop(0)

    def get_smoothed_metric(
        self, step_name: str, metric_name: str, window: int = 20
    ) -> Optional[float]:
        """
        Get smoothed metric value using moving average.

        Args:
            step_name: Name of step
            metric_name: Name of metric (e.g., 'vae_loss')
            window: Window size for moving average

        Returns:
            Smoothed metric value, or None if insufficient data
        """
        if step_name not in self.step_metrics:
            return None

        if metric_name not in self.step_metrics[step_name]:
            return None

        history = self.step_metrics[step_name][metric_name]
        if len(history) < window:
            return None

        return sum(history[-window:]) / window

    def should_transition(self, step: PipelineStepConfig, current_epoch: int) -> tuple[bool, str]:
        """
        Check if transition criteria are met for current step.

        Args:
            step: Current pipeline step configuration
            current_epoch: Current epoch number (within this step)

        Returns:
            Tuple of (should_transition, reason_string)
        """
        criteria = step.transition_on

        if criteria.mode == "epoch":
            if current_epoch >= step.n_epochs:
                return True, f"Completed {step.n_epochs} epochs"
            return False, f"Epoch {current_epoch}/{step.n_epochs}"

        elif criteria.mode == "loss_threshold":
            # Get smoothed metric value
            metric_value = self.get_smoothed_metric(step.name, criteria.metric)

            # Check max_epochs upper limit first
            max_epochs = criteria.max_epochs or step.n_epochs
            if current_epoch >= max_epochs:
                if metric_value is not None:
                    return (
                        True,
                        f"Max epochs ({max_epochs}) reached, {criteria.metric}={metric_value:.4f}",
                    )
                return True, f"Max epochs ({max_epochs}) reached"

            # Check if we have enough data for smoothed metric
            if metric_value is None:
                return False, f"Collecting metrics ({criteria.metric})"

            # Check threshold
            if metric_value < criteria.threshold:
                return (
                    True,
                    f"{criteria.metric}={metric_value:.4f} < {criteria.threshold} (threshold met)",
                )

            return (
                False,
                f"{criteria.metric}={metric_value:.4f} "
                f"(target: <{criteria.threshold}, epochs: {current_epoch}/{max_epochs})",
            )

        return False, "Unknown transition mode"

    def get_pipeline_metadata(self, step_index: int, step_epoch: int, batch_idx: int) -> dict:
        """
        Get pipeline metadata for checkpoint saving.

        Args:
            step_index: Current step index
            step_epoch: Current epoch within the current step (0-based)
            batch_idx: Current batch index within the current epoch

        Returns:
            Dictionary with pipeline state metadata
        """
        current_step = self.config.steps[step_index]

        return {
            "current_step_index": step_index,
            "current_step_name": current_step.name,
            "current_step_epoch": step_epoch,
            "current_batch_idx": batch_idx,
            "total_steps": len(self.config.steps),
            "steps_completed": self.steps_completed.copy(),
        }

    def resume_from_checkpoint(self) -> tuple[int, int, int]:
        """
        Resume pipeline from checkpoint if available.

        Returns:
            Tuple of (step_index, step_epoch, batch_idx)
                step_epoch: Epoch number within the current step (0-based)
        """
        training_state = self.checkpoint_manager.load_training_state()

        if not training_state:
            logger.info("No checkpoint found, starting from beginning")
            return 0, 0, 0

        # Check if this is a pipeline checkpoint
        if training_state.get("mode") != "pipeline":
            logger.info("Checkpoint is legacy mode (not pipeline), starting from step 0")
            return 0, 0, 0

        pipeline_meta = training_state.get("pipeline", {})
        step_index = pipeline_meta.get("current_step_index", 0)

        # Use step-local epoch from pipeline metadata (new format)
        # Fall back to global epoch for backward compatibility
        step_epoch = pipeline_meta.get("current_step_epoch", training_state.get("epoch", 0))

        # Use batch_idx from pipeline metadata if available, else from training state
        batch_idx = pipeline_meta.get("current_batch_idx", training_state.get("batch_idx", 0))

        self.global_step = training_state.get("global_step", 0)
        self.steps_completed = pipeline_meta.get("steps_completed", [])

        logger.info(
            f"Resuming from checkpoint: "
            f"step {step_index + 1}/{len(self.config.steps)} "
            f"('{pipeline_meta.get('current_step_name', 'unknown')}'), "
            f"step_epoch {step_epoch + 1}, batch {batch_idx}, global_step {self.global_step}"
        )

        return step_index, step_epoch, batch_idx

    def print_pipeline_summary(self) -> None:
        """Print pipeline execution plan summary."""
        print("\n" + "=" * 80)
        print("PIPELINE EXECUTION PLAN")
        print("=" * 80)

        total_epochs = sum(step.n_epochs for step in self.config.steps)

        for i, step in enumerate(self.config.steps, 1):
            print(f"\nStep {i}/{len(self.config.steps)}: {step.name} ({step.n_epochs} epochs)")

            if step.description:
                print(f"  Description: {step.description}")

            # Show training modes
            modes = []
            if step.train_vae:
                modes.append("VAE")
            if step.gan_training:
                modes.append("GAN")
            if step.train_spade:
                modes.append("SPADE")
            if step.use_lpips:
                modes.append("LPIPS")
            if step.train_diff or step.train_diff_full:
                modes.append("Flow")
            print(f"  Train: {', '.join(modes)}")

            # Show frozen models
            if step.freeze:
                print(f"  Frozen: {', '.join(step.freeze)}")

            # Show transition criteria
            if step.transition_on.mode == "epoch":
                print(f"  Transition: After {step.n_epochs} epochs")
            elif step.transition_on.mode == "loss_threshold":
                print(
                    f"  Transition: When {step.transition_on.metric} < "
                    f"{step.transition_on.threshold} (max {step.transition_on.max_epochs} epochs)"
                )

        print(f"\nTotal epochs: {total_epochs}")
        print("=" * 80 + "\n")

    def _create_step_optimizers(self, step, models, args):
        """Create optimizers for current step from inline config or defaults."""
        from ..training.optimizer_factory import create_optimizer

        optimizers = {}

        if not step.optimization or not step.optimization.optimizers:
            # Create default optimizers based on training modes
            logger.info("No optimizer config found, creating defaults based on training modes")

            # Default optimizer config
            default_opt_config = {
                "type": "AdamW",
                "lr": (
                    step.lr
                    if hasattr(step, "lr") and step.lr
                    else (args.lr if hasattr(args, "lr") else 0.0001)
                ),
                "weight_decay": 0.01,
                "eps": 1e-8,
                "betas": (0.9, 0.999),
            }

            # Create VAE optimizer if training VAE/GAN/SPADE/LPIPS
            needs_vae_trainer = (
                step.train_vae or step.gan_training or step.train_spade or step.use_lpips
            )
            if needs_vae_trainer:
                vae_params = list(models["compressor"].parameters()) + list(
                    models["expander"].parameters()
                )
                optimizers["vae"] = create_optimizer(vae_params, default_opt_config)
                logger.info(
                    f"✓ Created default VAE optimizer: AdamW (lr={default_opt_config['lr']:.2e})"
                )

            # Create discriminator optimizer if GAN enabled
            if step.gan_training and models.get("D_img"):
                optimizers["discriminator"] = create_optimizer(
                    models["D_img"].parameters(), default_opt_config
                )
                logger.info(
                    f"✓ Created default discriminator optimizer: AdamW (lr={default_opt_config['lr']:.2e})"
                )

            # Create flow optimizer if training flow
            if (step.train_diff or step.train_diff_full) and models.get("flow_processor"):
                optimizers["flow"] = create_optimizer(
                    models["flow_processor"].parameters(), default_opt_config
                )
                logger.info(
                    f"✓ Created default flow optimizer: AdamW (lr={default_opt_config['lr']:.2e})"
                )

            # Create text encoder optimizer if specified
            if (
                hasattr(step, "train_text_encoder")
                and step.train_text_encoder
                and models.get("text_encoder")
            ):
                optimizers["text_encoder"] = create_optimizer(
                    models["text_encoder"].parameters(), default_opt_config
                )
                logger.info(
                    f"✓ Created default text_encoder optimizer: AdamW (lr={default_opt_config['lr']:.2e})"
                )

            return optimizers

        # Explicit optimizer config provided - rest of method unchanged
        for name, opt_config_obj in step.optimization.optimizers.items():
            # Convert dataclass to dict, filtering out None values
            if hasattr(opt_config_obj, "__dataclass_fields__"):
                opt_config = {k: v for k, v in asdict(opt_config_obj).items() if v is not None}
            else:
                opt_config = opt_config_obj

            # Determine which parameters to optimize
            if name == "vae":
                params = list(models["compressor"].parameters()) + list(
                    models["expander"].parameters()
                )
            elif name == "discriminator":
                params = models["D_img"].parameters()
            elif name == "flow":
                params = models["flow_processor"].parameters()
            elif name == "text_encoder":
                params = models["text_encoder"].parameters()
            else:
                logger.warning(f"Unknown optimizer name: {name}, skipping")
                continue

            optimizers[name] = create_optimizer(params, opt_config)
            logger.info(f"Created optimizer for {name}: {opt_config.get('type', 'AdamW')}")

        return optimizers

    def _create_step_schedulers(self, step, optimizers, total_steps):
        """Create schedulers for current step from inline config."""
        from ..training.scheduler_factory import create_scheduler

        schedulers = {}

        if not step.optimization or not step.optimization.schedulers:
            logger.info("No scheduler config found, skipping")
            return schedulers

        for name, sched_config_obj in step.optimization.schedulers.items():
            if name not in optimizers:
                logger.warning(f"Scheduler '{name}' has no corresponding optimizer")
                continue

            # Convert SchedulerConfig dataclass to dict, filtering out None values
            if hasattr(sched_config_obj, "__dataclass_fields__"):
                sched_config = {k: v for k, v in asdict(sched_config_obj).items() if v is not None}
            else:
                sched_config = sched_config_obj

            scheduler = create_scheduler(optimizers[name], sched_config, total_steps)
            schedulers[name] = scheduler
            logger.info(f"Created scheduler '{name}': {sched_config['type']}")

        return schedulers

    def _create_step_trainers(self, step, models, optimizers, schedulers, ema, args):
        """Create trainers for current step."""
        import torch.nn as nn

        from ..training import FlowTrainer, VAETrainer

        trainers = {}

        # Create VAE trainer if training VAE, GAN, SPADE, or LPIPS
        # (these all require the VAE trainer even if train_vae=false)
        needs_vae_trainer = (
            step.train_vae or step.gan_training or step.train_spade or step.use_lpips
        )

        if needs_vae_trainer and "vae" in optimizers:
            trainers["vae"] = VAETrainer(
                compressor=models["compressor"],
                expander=models["expander"],
                optimizer=optimizers["vae"],
                scheduler=schedulers.get("vae"),
                ema=ema,
                reconstruction_loss_fn=nn.L1Loss(),
                reconstruction_loss_min_fn=nn.MSELoss(),
                use_spade=step.train_spade,
                train_reconstruction=step.train_vae,  # Only compute recon loss if train_vae=True
                kl_beta=step.kl_beta if hasattr(step, "kl_beta") else 0.0001,
                kl_warmup_steps=step.kl_warmup_steps if hasattr(step, "kl_warmup_steps") else 5000,
                kl_free_bits=step.kl_free_bits if hasattr(step, "kl_free_bits") else 0.0,
                use_gan=step.gan_training,
                discriminator=models["D_img"] if step.gan_training else None,
                discriminator_optimizer=optimizers.get("discriminator"),
                discriminator_scheduler=schedulers.get("discriminator"),
                lambda_adv=step.lambda_adv if hasattr(step, "lambda_adv") else 0.5,
                use_lpips=step.use_lpips,
                lambda_lpips=step.lambda_lpips if hasattr(step, "lambda_lpips") else 0.1,
                r1_gamma=5.0,
                r1_interval=16,
                gradient_clip_norm=args.initial_clipping_norm,
                accelerator=self.accelerator,
            )

            # Build descriptive message about what's being trained
            modes = []
            if step.train_vae:
                modes.append("VAE")
            if step.train_spade:
                modes.append("SPADE")
            if step.gan_training:
                modes.append("GAN")
            if step.use_lpips:
                modes.append("LPIPS")

            logger.info(f"Created VAE trainer ({', '.join(modes)})")

        if (step.train_diff or step.train_diff_full) and "flow" in optimizers:
            trainers["flow"] = FlowTrainer(
                flow_processor=models["flow_processor"],
                text_encoder=models["text_encoder"],
                compressor=models["compressor"],
                optimizer=optimizers["flow"],
                scheduler=schedulers.get("flow"),
                text_encoder_optimizer=optimizers.get("text_encoder"),
                text_encoder_scheduler=schedulers.get("text_encoder"),
                gradient_clip_norm=args.initial_clipping_norm,
                num_train_timesteps=1000,
                accelerator=self.accelerator,
            )
            logger.info("Created Flow trainer")

        return trainers

    def _save_checkpoint(
        self, step_idx, step_epoch, batch_idx, models, optimizers, schedulers, ema, args
    ):
        """
        Save checkpoint with pipeline metadata.

        Args:
            step_idx: Current pipeline step index
            step_epoch: Current epoch within the step (0-based)
            batch_idx: Current batch index
            models: Dictionary of models
            optimizers: Dictionary of optimizers
            schedulers: Dictionary of schedulers
            ema: EMA module (if applicable)
            args: Training arguments
        """
        # Get pipeline metadata
        metadata = self.get_pipeline_metadata(step_idx, step_epoch, batch_idx)

        # Save models
        self.checkpoint_manager.save_models(
            diffuser=models["diffuser"],
            text_encoder=models["text_encoder"],
            discriminators={"D_img": models["D_img"]} if models.get("D_img") else None,
        )

        # Save training state with pipeline metadata
        self.checkpoint_manager.save_training_state(
            epoch=step_epoch,  # Use step-local epoch for consistency
            batch_idx=batch_idx,
            global_step=self.global_step,
            optimizers=optimizers,
            schedulers=schedulers,
            ema=ema,
            pipeline_metadata=metadata,
        )

        logger.info(
            f"Checkpoint saved: step {step_idx+1}/{len(self.config.steps)}, "
            f"epoch {step_epoch+1}, batch {batch_idx}, global_step {self.global_step}"
        )

    def _generate_samples(
        self, step, step_idx, epoch, batch_idx, models, tokenizer, args, parsed_sample_sizes
    ):
        """
        Generate sample images for monitoring training progress.

        Args:
            step: Current pipeline step config
            step_idx: Current step index
            epoch: Current epoch within step
            batch_idx: Current batch index
            models: Dictionary of models
            tokenizer: Tokenizer instance
            args: Training arguments
            parsed_sample_sizes: List of sample size tuples
        """
        import glob
        import os

        if args.no_samples:
            return

        diffuser = models.get("diffuser")
        text_encoder = models.get("text_encoder")
        if not diffuser:
            logger.warning("Cannot generate samples: diffuser model not found")
            return

        # Sample epoch identifier (use global_step for legacy compatibility)
        sample_epoch = self.global_step

        logger.info(
            f"Generating samples for step {step_idx+1}/{len(self.config.steps)}, "
            f"epoch {epoch+1}, batch {batch_idx}, global_step {sample_epoch}"
        )

        # Create prefix for new naming: stepname_step_epoch_batch
        step_name_short = step.name[:20]  # Limit step name length
        # Include batch for mid-epoch samples, omit for end-of-epoch (batch=-1 or max batch)
        if batch_idx >= 0 and batch_idx < 999999:  # Mid-epoch
            sample_prefix = f"{step_name_short}_{step_idx+1:03d}_{epoch+1:03d}_{batch_idx:05d}"
        else:  # End-of-epoch or initial
            sample_prefix = f"{step_name_short}_{step_idx+1:03d}_{epoch+1:03d}"

        # VAE reconstruction samples (if test images provided)
        if args.test_image_address and len(args.test_image_address) > 0:
            for img_addr in args.test_image_address:
                try:
                    # Generate samples with epoch number (legacy naming)
                    safe_vae_sample(
                        diffuser,
                        img_addr,
                        args.channels if hasattr(args, "channels") else 3,
                        args.output_path,
                        sample_epoch,
                        self.device,
                    )

                    # Rename files to include step/epoch info
                    # Pattern: vae_epoch_0031-{hash}-{suffix}.webp → {prefix}_{hash}-{suffix}.webp
                    pattern = os.path.join(args.output_path, f"vae_epoch_{sample_epoch:04d}-*.webp")
                    for old_file in glob.glob(pattern):
                        old_filename = os.path.basename(old_file)
                        # Extract hash and suffix: "vae_epoch_0031-cd7a10...-ctx.webp"
                        # Split on first '-' after epoch number
                        parts = old_filename.replace(f"vae_epoch_{sample_epoch:04d}-", "")
                        # parts = "cd7a10...-ctx.webp"
                        hash_suffix = parts.replace(".webp", "")

                        # New filename: stepname_step_epoch_hash-suffix.webp
                        new_filename = f"{sample_prefix}_{hash_suffix}.webp"
                        new_file = os.path.join(args.output_path, new_filename)

                        os.rename(old_file, new_file)
                        logger.debug(f"Renamed: {old_filename} → {new_filename}")

                except Exception as e:
                    logger.warning(f"Failed to generate VAE sample from {img_addr}: {e}")

        # Flow-based text-to-image samples (if flow training active)
        if (step.train_diff or step.train_diff_full) and text_encoder:
            if args.sample_captions and len(args.sample_captions) > 0:
                try:
                    # Generate samples with epoch number (legacy naming)
                    save_sample_images(
                        diffuser,
                        text_encoder,
                        tokenizer,
                        args.output_path,
                        sample_epoch,
                        self.device,
                        args.sample_captions,
                        args.batch_size,
                        sample_sizes=parsed_sample_sizes,
                    )

                    # Rename files to include step/epoch/batch info
                    # Pattern: sample_epoch_<N>_<caption_idx>_<size>.png
                    pattern = os.path.join(args.output_path, f"sample_epoch_{sample_epoch}_*.png")
                    for old_file in glob.glob(pattern):
                        old_filename = os.path.basename(old_file)
                        # Extract caption index and size: "sample_epoch_247_0_512x512.png"
                        parts = old_filename.replace(f"sample_epoch_{sample_epoch}_", "").replace(
                            ".png", ""
                        )
                        # parts = "0_512x512" or just "0"

                        # New filename: stepname_step_epoch_batch_captionidx_size.png
                        new_filename = f"{sample_prefix}_{parts}.png"
                        new_file = os.path.join(args.output_path, new_filename)

                        os.rename(old_file, new_file)
                        logger.debug(f"Renamed: {old_filename} → {new_filename}")

                except Exception as e:
                    logger.warning(f"Failed to generate flow samples: {e}")

    def run(self, models, dataloader, sampler, tokenizer, progress_logger, args, config) -> None:
        """
        Execute the complete training pipeline.

        This method is the main entry point for pipeline execution. It orchestrates
        multi-step training by configuring models, creating trainers, and managing
        the training loop across pipeline steps.

        Args:
            models: Dict of initialized models:
                - diffuser: FluxPipeline instance
                - compressor: FluxCompressor instance
                - expander: FluxExpander instance
                - flow_processor: FluxFlowProcessor instance
                - text_encoder: BertTextEncoder instance
                - D_img: PatchDiscriminator instance (if GAN training)
            dataloader: Initialized DataLoader for training data
            tokenizer: Tokenizer for text processing
            args: Parsed command-line arguments
            config: Loaded YAML config dictionary

        Raises:
            NotImplementedError: Full implementation deferred to Phase 3b
                                See docs/PIPELINE_ARCHITECTURE.md for design

        Architecture Overview:
            1. Resume from checkpoint (if exists) → get start_step, start_epoch, start_batch
            2. For each step in pipeline:
                a. configure_step_models() - freeze/unfreeze per step config
                b. _create_step_optimizers() - from inline YAML config
                c. _create_step_schedulers() - from inline YAML config
                d. _create_step_trainers() - VAETrainer and/or FlowTrainer
                e. Training loop:
                   - For each epoch in step:
                     - For each batch:
                       - vae_trainer.train_step() if train_vae
                       - flow_trainer.train_step() if train_flow
                       - update_metrics()
                       - log progress
                       - save checkpoint (with pipeline metadata)
                     - Check transition_criteria (epoch or loss_threshold)
                f. Cleanup optimizers/schedulers
            3. Print final summary

        Next Implementation Steps:
            1. Extract initialize_models() and initialize_dataloader() helpers from train_legacy()
            2. Implement _create_step_optimizers() using create_optimizer() factory
            3. Implement _create_step_schedulers() using create_scheduler() factory
            4. Implement _create_step_trainers() using VAETrainer/FlowTrainer
            5. Implement main training loop with transition monitoring
            6. Implement _save_checkpoint() with pipeline metadata
            7. Add integration tests

        For detailed architecture and implementation plan:
            See docs/PIPELINE_ARCHITECTURE.md
        """
        import time

        import torch
        import torch.nn as nn
        from fluxflow.utils import format_duration

        from ..training import EMA, FloatBuffer

        logger.info("Starting training pipeline execution...")

        # Print pipeline summary
        self.print_pipeline_summary()

        # Resume from checkpoint if available
        start_step, start_epoch, start_batch = self.resume_from_checkpoint()

        logger.info(
            f"Pipeline has {len(self.config.steps)} steps, starting from step {start_step + 1}"
        )

        # Parse sample sizes for sample generation
        from ..scripts.train import parse_sample_sizes

        parsed_sample_sizes = parse_sample_sizes(
            config.get("output", {}).get("sample_sizes", [512])
        )

        # Get dataset size for progress tracking
        if isinstance(dataloader.dataset, torch.utils.data.IterableDataset):
            dataset_size = getattr(dataloader.dataset, "dataset_size", 1000)
        else:
            dataset_size = len(dataloader.dataset)

        batches_per_epoch = max(1, dataset_size // args.batch_size)

        # Generate initial samples (before training starts)
        if start_step == 0 and start_epoch == 0:
            logger.info("Generating initial samples before training...")
            self._generate_samples(
                self.config.steps[0], 0, -1, 0, models, tokenizer, args, parsed_sample_sizes
            )

        # Main pipeline loop
        for step_idx in range(start_step, len(self.config.steps)):
            step = self.config.steps[step_idx]

            print(f"\n{'='*80}")
            print(f"PIPELINE STEP {step_idx+1}/{len(self.config.steps)}: {step.name}")
            if step.description:
                print(f"Description: {step.description}")
            print(f"Duration: {step.n_epochs} epochs")
            print(f"{'='*80}\n")

            # Configure models for this step (freeze/unfreeze)
            self.configure_step_models(step, models)

            # Update progress logger for this step (step-specific files)
            progress_logger.set_step(step.name)
            logger.info(f"Progress logging to: {progress_logger.metrics_file}")

            # Create optimizers and schedulers for this step
            optimizers = self._create_step_optimizers(step, models, args)
            total_steps = step.n_epochs * batches_per_epoch
            schedulers = self._create_step_schedulers(step, optimizers, total_steps)

            # Create EMA if training VAE
            # Create EMA if we need VAE trainer (for VAE, GAN, SPADE, or LPIPS)
            needs_vae_trainer = step.train_vae or step.train_spade or step.use_lpips
            ema = None
            if needs_vae_trainer and step.use_ema:
                ema = EMA(
                    nn.ModuleList([models["compressor"], models["expander"]]),
                    decay=0.999,
                    device=self.device,
                )
                logger.info("✓ EMA enabled (adds 2x model VRAM)")
            elif needs_vae_trainer and not step.use_ema:
                logger.info("⚠ EMA disabled to save VRAM (~14GB for vae_dim=128)")

            # Create trainers for this step
            trainers = self._create_step_trainers(step, models, optimizers, schedulers, ema, args)

            # Training loop for this step
            step_start_time = time.time()

            for epoch in range(start_epoch if step_idx == start_step else 0, step.n_epochs):
                # Calculate total batches for this epoch (considering max_steps)
                epoch_total_batches = (
                    min(batches_per_epoch, step.max_steps) if step.max_steps else batches_per_epoch
                )

                print(
                    f"\nStep {step.name} ({step_idx+1}/{len(self.config.steps)}), "
                    f"Epoch {epoch+1}/{step.n_epochs}, "
                    f"Batches 0/{epoch_total_batches}"
                )

                epoch_start_time = time.time()

                # Error buffers for logging
                vae_errors = FloatBuffer(max(args.log_interval * 2, 10))
                kl_errors = FloatBuffer(max(args.log_interval * 2, 10))
                flow_errors = FloatBuffer(max(args.log_interval * 2, 10))
                g_errors = FloatBuffer(max(args.log_interval * 2, 10))  # GAN generator loss
                d_errors = FloatBuffer(max(args.log_interval * 2, 10))  # GAN discriminator loss
                lpips_errors = FloatBuffer(max(args.log_interval * 2, 10))  # LPIPS loss
                batch_times = FloatBuffer(max(args.log_interval * 2, 10))  # Batch timing

                for batch_idx, (imgs, input_ids) in enumerate(dataloader):
                    batch_start_time = time.time()
                    # Break if max_steps reached (for quick testing)
                    if step.max_steps is not None and batch_idx >= step.max_steps:
                        logger.info(f"Reached max_steps={step.max_steps}, ending epoch early")
                        break

                    # Skip batches if resuming mid-epoch
                    if step_idx == start_step and epoch == start_epoch and batch_idx < start_batch:
                        continue

                    self.global_step += 1
                    input_ids = input_ids.to(self.device)
                    attention_mask = (input_ids != tokenizer.pad_token_id).long().to(self.device)

                    # Train on all resolutions
                    for ri in imgs:
                        real_imgs = ri.to(self.device).detach()

                        # VAE/GAN/SPADE training (runs if trainer exists, even with train_vae=false)
                        if trainers.get("vae"):
                            vae_losses = trainers["vae"].train_step(real_imgs, self.global_step)
                            vae_errors.add_item(vae_losses["vae"])
                            kl_errors.add_item(vae_losses["kl"])

                            # Track GAN losses if available
                            if "generator" in vae_losses:
                                g_errors.add_item(vae_losses["generator"])
                            if "discriminator" in vae_losses:
                                d_errors.add_item(vae_losses["discriminator"])
                            if "lpips" in vae_losses:
                                lpips_errors.add_item(vae_losses["lpips"])

                            # Update metrics for transition monitoring
                            self.update_metrics(step.name, {"vae_loss": vae_losses["vae"]})

                        # Flow training
                        if (step.train_diff or step.train_diff_full) and trainers.get("flow"):
                            flow_losses = trainers["flow"].train_step(
                                real_imgs, input_ids, attention_mask
                            )
                            flow_loss = (
                                flow_losses["flow_loss"]
                                if isinstance(flow_losses, dict)
                                else flow_losses
                            )
                            flow_errors.add_item(flow_loss)

                            # Update metrics for transition monitoring
                            self.update_metrics(step.name, {"flow_loss": flow_loss})

                    # Track batch time
                    batch_time = time.time() - batch_start_time
                    batch_times.add_item(batch_time)

                    # Periodic CUDA cache clearing to prevent fragmentation (every 10 batches)
                    if batch_idx % 10 == 0:
                        torch.cuda.empty_cache()

                    # Logging
                    if batch_idx % args.log_interval == 0:
                        elapsed = time.time() - step_start_time
                        elapsed_str = format_duration(elapsed)

                        log_msg = (
                            f"[{elapsed_str}] Step {step.name} ({step_idx+1}/{len(self.config.steps)}) | "
                            f"Epoch {epoch+1}/{step.n_epochs} | "
                            f"Batch {batch_idx}/{epoch_total_batches}"
                        )

                        # Console logging (show metrics if VAE trainer ran)
                        if len(vae_errors._items) > 0:
                            log_msg += (
                                f" | VAE: {vae_errors.average:.4f} | KL: {kl_errors.average:.4f}"
                            )
                            # Add GAN losses if active
                            if step.gan_training and len(g_errors._items) > 0:
                                log_msg += (
                                    f" | G: {g_errors.average:.4f} | D: {d_errors.average:.4f}"
                                )
                            # Add LPIPS if active
                            if step.use_lpips and len(lpips_errors._items) > 0:
                                log_msg += f" | LPIPS: {lpips_errors.average:.4f}"

                        if step.train_diff or step.train_diff_full:
                            log_msg += f" | Flow: {flow_errors.average:.4f}"

                        # Add average batch time
                        if len(batch_times._items) > 0:
                            log_msg += f" | {batch_times.average:.2f}s/batch"

                        print(log_msg)

                        # Log to progress logger (include metrics if VAE trainer ran)
                        metrics = {}
                        if len(vae_errors._items) > 0:
                            metrics["vae_loss"] = vae_errors.average
                            metrics["kl_loss"] = kl_errors.average
                            # Add GAN metrics
                            if step.gan_training and len(g_errors._items) > 0:
                                metrics["g_loss"] = g_errors.average
                                metrics["d_loss"] = d_errors.average
                            # Add LPIPS metrics
                            if step.use_lpips and len(lpips_errors._items) > 0:
                                metrics["lpips_loss"] = lpips_errors.average

                        if step.train_diff or step.train_diff_full:
                            metrics["flow_loss"] = flow_errors.average

                        progress_logger.log_metrics(
                            epoch=epoch,
                            batch=batch_idx,
                            global_step=self.global_step,
                            metrics=metrics,
                            learning_rates={},
                        )

                    # Checkpoint saving (mid-epoch)
                    if batch_idx % args.checkpoint_save_interval == 0 and batch_idx > 0:
                        self._save_checkpoint(
                            step_idx, epoch, batch_idx, models, optimizers, schedulers, ema, args
                        )

                        # Clear CUDA cache after checkpoint save to prevent fragmentation
                        torch.cuda.empty_cache()

                        # Generate samples at checkpoint intervals if requested
                        if args.samples_per_checkpoint > 0:
                            self._generate_samples(
                                step,
                                step_idx,
                                epoch,
                                batch_idx,
                                models,
                                tokenizer,
                                args,
                                parsed_sample_sizes,
                            )

                # End-of-epoch checkpoint (always save after completing an epoch)
                epoch_time = time.time() - epoch_start_time
                print(f"Epoch {epoch+1} completed in {format_duration(epoch_time)}")

                self._save_checkpoint(
                    step_idx, epoch, batch_idx, models, optimizers, schedulers, ema, args
                )
                logger.info(f"End-of-epoch checkpoint saved")

                # Generate samples after epoch completes (use last batch_idx)
                self._generate_samples(
                    step, step_idx, epoch, batch_idx, models, tokenizer, args, parsed_sample_sizes
                )

                # Check transition criteria (after saving checkpoint)
                should_trans, reason = self.should_transition(step, epoch)
                if should_trans:
                    print(f"\nTransition criteria met: {reason}")
                    print(f"Moving to next step...")
                    # Save checkpoint before transitioning
                    self._save_checkpoint(
                        step_idx, epoch, batch_idx, models, optimizers, schedulers, ema, args
                    )
                    logger.info("Pre-transition checkpoint saved")
                    break

            # Save final checkpoint at end of step
            logger.info(f"Step {step_idx+1} complete, saving final checkpoint")
            self._save_checkpoint(
                step_idx, epoch, batch_idx, models, optimizers, schedulers, ema, args
            )

            # Mark step as completed
            self.steps_completed.append(step.name)

            # Cleanup
            del optimizers, schedulers, trainers
            if ema:
                del ema
            torch.cuda.empty_cache()

            # Reset start_epoch and start_batch for next step
            start_epoch = 0
            start_batch = 0

        print(f"\n{'='*80}")
        print("PIPELINE TRAINING COMPLETE")
        print(f"{'='*80}\n")
