#!/usr/bin/env python3
"""
Simple inference script for Flux.Kontext with optional RL LoRA support.
Supports both standard LoRA and RL LoRA inference modes.
"""

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"  # Suppress TensorFlow logs
import argparse
from pathlib import Path
import torch
from diffusers import PEFluxKontextPipeline
from diffusers.utils import load_image
from safetensors.torch import load_file


def initialize_pipeline(
    base_model_path: str,
    rl_lora_path: str = None,
    device: str = "cuda"
):
    """
    Initialize Flux.Kontext pipeline with optional LoRA weights.
    
    Args:
        base_model_path: Path to base Flux model
        lora_path: Path to standard LoRA weights (optional)
        rl_lora_path: Path to RL LoRA weights (optional)
        device: Device to run inference on
    
    Returns:
        Initialized pipeline
    """
    print(f"Initializing Flux.Kontext pipeline on {device}...")
    
    # Load base pipeline
    pipe = PEFluxKontextPipeline.from_pretrained(
        base_model_path,
        torch_dtype=torch.bfloat16,
    ).to(device)
    
    # Load RL LoRA if provided
    if rl_lora_path:
        print(f"Loading RL LoRA from: {rl_lora_path}")
        rl_sd = load_file(rl_lora_path)
        
        # Convert keys from base_model.model to transformer
        new_rl_sd = {}
        for key in rl_sd.keys():
            new_key = key.replace("base_model.model.", "transformer.")
            new_rl_sd[new_key] = rl_sd[key]
        
        pipe.load_lora_weights(new_rl_sd, adapter_name='rl_lora')
    
    print("Pipeline initialized successfully!")
    return pipe


def run_inference(
    pipe,
    input_image_path: str,
    reference_image_path: str,
    prompt: str,
    output_path: str,
    use_rl_lora: bool = False,
    num_inference_steps: int = 28,
    guidance_scale: float = 3.5,
    height: int = None,
    width: int = None,
    seed: int = 42
):
    """
    Run inference with the pipeline.
    
    Args:
        pipe: Initialized pipeline
        input_image_path: Path to input image to fix
        reference_image_path: Path to reference image
        prompt: Text prompt for inference
        output_path: Path to save output image
        use_rl_lora: Whether to use RL LoRA (if False, uses standard LoRA or no LoRA)
        num_inference_steps: Number of diffusion steps
        guidance_scale: Guidance scale for inference
        height: Target height (auto-calculated if None)
        width: Target width (auto-calculated if None)
        seed: Random seed for reproducibility
    
    Returns:
        Generated image
    """
    print(f"\nRunning inference...")
    print(f"  Mode: {'RL LoRA' if use_rl_lora else 'Standard LoRA / Base model'}")
    print(f"  Input: {input_image_path}")
    print(f"  Reference: {reference_image_path}")
    print(f"  Prompt: {prompt}")
    
    # Load images
    input_image = load_image(input_image_path)
    reference_image = load_image(reference_image_path)
    
    # Auto-calculate dimensions if not provided
    if height is None or width is None:
        target_height, target_width = input_image.height, input_image.width
        
        # Scale to appropriate resolution
        if target_height < 512:
            scale_factor = 512 / target_height
            target_height = 512
            target_width = int(target_width * scale_factor)
        elif target_height > 512:
            scale_factor = 1024 / target_height
            target_height = 1024
            target_width = int(target_width * scale_factor)
        
        height = target_height
        width = target_width
    
    print(f"  Resolution: {width}x{height}")
    
    # Set LoRA mode
    if use_rl_lora:
        pipe.enable_lora()
        # Uncomment if you need to explicitly set adapter
        # pipe.set_adapters('rl_lora', 1.0)
    else:
        # For standard LoRA or base model
        pipe.disable_lora()
        # Uncomment if you loaded standard LoRA and want to use it
        # pipe.set_adapters('lora', 1.0)
    
    # Run inference
    result = pipe(
        prompt=prompt,
        image=input_image,
        reference=reference_image,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        height=height,
        width=width,
        generator=torch.Generator(device=pipe.device).manual_seed(seed)
    ).images[0]
    
    # Save result
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    result.save(output_path)
    
    print(f"  ✓ Saved to: {output_path}")
    
    return result


def main():
    parser = argparse.ArgumentParser(
        description="Simple inference script for Flux.Kontext with optional RL LoRA"
    )
    
    # Model paths
    parser.add_argument(
        "--base_model",
        type=str,
        default="./models/base",
        # default="/raid/users/lyl/DanceGRPO-kontext/data/refinerany_20250929",
        help="Path to base Flux model"
    )
    parser.add_argument(
        "--rl_lora_path",
        type=str,
        # default="./models/adapter_model.safetensors",
        default=None,
        help="Path to RL LoRA weights (optional)"
    )
    
    # Input/output
    parser.add_argument(
        "--input_image",
        type=str,
        default="./input.png",
        help="Path to input image to fix"
    )
    parser.add_argument(
        "--reference_image",
        type=str,
        default="./reference.png",
        help="Path to reference image"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="output/result.png",
        help="Path to save output image"
    )
    parser.add_argument(
        "--prompt",
        type=str,
        default="Fix the control image according to the referenced image: Use the wheels from the referenced image to repair the front and rear wheels of the car in the control image. Make sure to preserve the Union Jack-style rims.",
        help="Text prompt for inference"
    )
    
    parser.add_argument(
        "--steps",
        type=int,
        default=28,
        help="Number of inference steps"
    )
    parser.add_argument(
        "--guidance_scale",
        type=float,
        default=3.5,
        help="Guidance scale"
    )
    parser.add_argument(
        "--height",
        type=int,
        default=None,
        help="Target height (auto-calculated if not specified)"
    )
    parser.add_argument(
        "--width",
        type=int,
        default=None,
        help="Target width (auto-calculated if not specified)"
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for reproducibility"
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        help="Device to run inference on"
    )
    parser.add_argument(
        "--gpu_id",
        type=str,
        default=None,
        help="GPU ID to use (e.g., '0' or '1')"
    )
    
    args = parser.parse_args()
    
    # Set GPU if specified
    if args.gpu_id is not None:
        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
        print(f"Using GPU: {args.gpu_id}")
    
    # Validate inputs
    if not Path(args.input_image).exists():
        print(f"Error: Input image not found: {args.input_image}")
        return
    
    if not Path(args.reference_image).exists():
        print(f"Error: Reference image not found: {args.reference_image}")
        return

    
    # Initialize pipeline
    try:
        pipe = initialize_pipeline(
            base_model_path=args.base_model,
            rl_lora_path=args.rl_lora_path,
            device=args.device
        )
    except Exception as e:
        print(f"Error initializing pipeline: {e}")
        import traceback
        traceback.print_exc()
        return
    
    # Run inference
    try:
        run_inference(
            pipe=pipe,
            input_image_path=args.input_image,
            reference_image_path=args.reference_image,
            prompt=args.prompt,
            output_path=args.output,
            use_rl_lora=args.rl_lora_path is not None,
            num_inference_steps=args.steps,
            guidance_scale=args.guidance_scale,
            height=args.height,
            width=args.width,
            seed=args.seed
        )
        print("\n✓ Inference completed successfully!")
    except Exception as e:
        print(f"✗ Error during inference: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()