"""Load the foveated FLUX2 pipeline plus optional LoRA / DiT checkpoints."""

import torch

from diffsynth.core import load_state_dict
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig

from ..diffsynth_fov import Flux2FoveatedImagePipeline


def load_pipeline(args, use_foveated_pipeline: bool = True):
    """Load FLUX2 (foveated by default) and optionally apply LoRA / replace DiT weights.

    For the user-study experiment LoRA / DiT swaps happen later (inside the runner)
    so the baseline + naive passes can use the base DiT.
    """
    print("Loading FLUX2 foveated pipeline...")
    pipeline_class = Flux2FoveatedImagePipeline if use_foveated_pipeline else Flux2ImagePipeline
    pipe = pipeline_class.from_pretrained(
        torch_dtype=torch.bfloat16,
        device="cuda" if torch.cuda.is_available() else "cpu",
        model_configs=[
            ModelConfig(model_id=args.model_id, origin_file_pattern="transformer/*.safetensors"),
            ModelConfig(model_id=args.model_id, origin_file_pattern="text_encoder/*.safetensors"),
            ModelConfig(model_id=args.model_id, origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
        ],
        tokenizer_config=ModelConfig(model_id=args.model_id, origin_file_pattern="tokenizer/"),
    )

    defer_load = (args.experiment == "user_study")

    if args.lora_checkpoint is not None and not defer_load:
        pipe.load_lora(pipe.dit, args.lora_checkpoint)
        print(f"Loaded LoRA checkpoint from {args.lora_checkpoint}")
    elif args.lora_checkpoint is not None and defer_load:
        print("User study: LoRA will be loaded only for 'ours' runs")

    if args.dit_checkpoint is not None and not defer_load:
        state_dict = load_state_dict(args.dit_checkpoint, torch_dtype=torch.bfloat16)
        pipe.dit.load_state_dict(state_dict)
        print(f"Loaded DiT checkpoint from {args.dit_checkpoint}")
    elif args.dit_checkpoint is not None and defer_load:
        print("User study: DiT will be loaded only for 'ours' runs")

    return pipe