#!/usr/bin/env python3
"""Compare Fine-Tuned vs Merged/Base+LoRA inference outputs.

Generates two videos with the same seed and computes SSIM.

Usage examples:
python lora_inference_comparison.py \
  --base ./merged_model \
  --ft FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers \
  --adapter NONE \
  --output-dir ./inference_comparison \
  --compute-ssim \
  --seed 41

or

python lora_inference_comparison.py \
  --base Wan-AI/Wan2.2-TI2V-5B-Diffusers \
  --ft FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers \
  --adapter adapter.safetensors \
  --output-dir ./inference_comparison \
  --compute-ssim \
  --seed 41

"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
from typing import Optional, Dict, Any

import logging

# minimal distributed env defaults (kept for compatibility)
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29500")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("LOCAL_RANK", "0")

# allow running from repo root where fastvideo is located
_FASTVIDEO_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "fastvideo_pr", "FastVideo"))
if _FASTVIDEO_PATH not in sys.path:
    sys.path.insert(0, _FASTVIDEO_PATH)

logger = logging.getLogger("inference_comparison")


def configure_logging(level: str = "INFO") -> None:
    handler = logging.StreamHandler()
    fmt = "%(asctime)s %(levelname)s %(message)s"
    handler.setFormatter(logging.Formatter(fmt, datefmt="%Y-%m-%d %H:%M:%S"))
    logger.addHandler(handler)
    logger.setLevel(level)


def _validate_adapter(adapter: Optional[str]) -> Optional[str]:
    if not adapter:
        return None
    if adapter.upper() == "NONE":
        return None
    p = Path(adapter).expanduser()
    if not p.exists():
        raise FileNotFoundError(f"Adapter not found: {p}")
    
    # Accept both files and directories (FastVideo expects directories for HF-style adapters)
    if p.is_file():
        if p.suffix != ".safetensors":
            raise ValueError(f"Adapter file must be .safetensors, got: {p.suffix}")
        if p.stat().st_size == 0:
            raise ValueError(f"Adapter file is empty: {p}")
    elif p.is_dir():
        # Check if directory contains at least one .safetensors file
        safetensors_files = list(p.glob("*.safetensors"))
        if not safetensors_files:
            raise ValueError(f"Adapter directory contains no .safetensors files: {p}")
    else:
        raise ValueError(f"Adapter must be a file or directory: {p}")
    
    return str(p.resolve())


def generate_with_model(
    model_path: str,
    output_dir: str,
    output_name: str,
    prompt: str,
    seed: int,
    lora_path: Optional[str],
    height: int,
    width: int,
    num_frames: int,
    num_inference_steps: int,
    guidance_scale: float,
    flow_shift: Optional[float] = None,
    embedded_guidance_scale: Optional[float] = None,
) -> str:
    """Produce a video with VideoGenerator.from_pretrained; returns video path."""
    try:
        from fastvideo import VideoGenerator  # lazy import
    except Exception as exc:
        raise RuntimeError(f"Failed to import fastvideo.VideoGenerator: {exc}") from exc

    init_kwargs: Dict[str, Any] = {
        "num_gpus": 1,
        "dit_cpu_offload": True,
        "vae_cpu_offload": True,
        "text_encoder_cpu_offload": True,
        "pin_cpu_memory": True,
    }
    if lora_path:
        init_kwargs["lora_path"] = lora_path
        init_kwargs["lora_nickname"] = "extracted"

    generator = VideoGenerator.from_pretrained(model_path, **init_kwargs)

    gen_kwargs = {
        "height": height,
        "width": width,
        "num_frames": num_frames,
        "num_inference_steps": num_inference_steps,
        "guidance_scale": guidance_scale,
        "seed": seed,
        "output_path": output_dir,
        "output_video_name": output_name,
        "save_video": True,
    }
    if flow_shift is not None:
        gen_kwargs["flow_shift"] = flow_shift
    if embedded_guidance_scale is not None:
        gen_kwargs["embedded_guidance_scale"] = embedded_guidance_scale

    result = generator.generate_video(prompt, **gen_kwargs)

    # best-effort cleanup of internal executors
    try:
        if hasattr(generator, "executor") and hasattr(generator.executor, "shutdown"):
            generator.executor.shutdown()
    except Exception:
        pass

    # determine saved video path
    expected = Path(output_dir) / f"{output_name}.mp4"
    if expected.exists():
        return str(expected)
    # fallback: check result dict
    if isinstance(result, dict) and "video_path" in result:
        return str(result["video_path"])
    raise FileNotFoundError(f"Video not found at expected path: {expected}")


def compute_metrics(output_dir: str, ft_video: str, other_video: str, num_inference_steps: int, prompt: str, compute_ssim: bool, compute_lpips: bool) -> dict:
    results = {}
    
    if compute_ssim:
        try:
            from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results  # type: ignore
            ssim_values = compute_video_ssim_torchvision(ft_video, other_video, use_ms_ssim=True)
            results["mean_ssim"] = float(ssim_values[0])
            write_ssim_results(output_dir, ssim_values, ft_video, other_video, num_inference_steps, prompt)
        except Exception as e:
            logger.warning(f"SSIM computation failed: {e}")
    
    if compute_lpips:
        try:
            import torch
            import lpips
            import torchvision.io as tv_io
            
            loss_fn = lpips.LPIPS(net='alex')
            
            # Load videos
            vid1, _, _ = tv_io.read_video(ft_video, pts_unit='sec')
            vid2, _, _ = tv_io.read_video(other_video, pts_unit='sec')
            
            # Normalize to [-1, 1]
            vid1 = (vid1.float() / 127.5 - 1.0).permute(0, 3, 1, 2)  # (T, C, H, W)
            vid2 = (vid2.float() / 127.5 - 1.0).permute(0, 3, 1, 2)
            
            lpips_scores = []
            with torch.no_grad():
                for frame1, frame2 in zip(vid1, vid2):
                    score = loss_fn(frame1.unsqueeze(0), frame2.unsqueeze(0))
                    lpips_scores.append(float(score.item()))
            
            results["mean_lpips"] = sum(lpips_scores) / len(lpips_scores)
            
            # Write LPIPS results
            import json
            lpips_file = Path(output_dir) / f"steps{num_inference_steps}_{prompt.replace(' ', '_')[:30]}_lpips.json"
            with open(lpips_file, 'w') as f:
                json.dump({"mean_lpips": results["mean_lpips"], "lpips_per_frame": lpips_scores}, f, indent=2)
                
        except Exception as e:
            logger.warning(f"LPIPS computation failed: {e}")
    
    return results


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Compare Fine-Tuned vs Merged/Base+LoRA inference outputs")
    p.add_argument("--base", required=True, help="Base model ID or merged model path")
    p.add_argument("--ft", required=True, help="Fine-tuned model ID or path (reference)")
    p.add_argument("--adapter", default="NONE", help="Path to .safetensors adapter, or NONE to use merged model")
    p.add_argument("--output-dir", default="./inference_comparison")
    p.add_argument("--prompt", default="A cat sitting on a windowsill")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--height", type=int, default=480)
    p.add_argument("--width", type=int, default=832)
    p.add_argument("--num-frames", type=int, default=49)
    p.add_argument("--num-inference-steps", type=int, default=32)
    p.add_argument("--guidance-scale", type=float, default=6.0)
    p.add_argument("--compute-ssim", action="store_true")
    p.add_argument("--compute-lpips", action="store_true")
    p.add_argument("--flow-shift", type=float, default=None)
    p.add_argument("--embedded-guidance-scale", type=float, default=None)
    p.add_argument("--log-level", default="INFO")
    return p.parse_args()


def compare_inference(
    base: str,
    ft: str,
    adapter: Optional[str],
    output_dir: str,
    prompt: str = "A cat sitting on a windowsill",
    seed: int = 42,
    height: int = 480,
    width: int = 832,
    num_frames: int = 49,
    num_inference_steps: int = 32,
    guidance_scale: float = 5.0,
    flow_shift: Optional[float] = None,
    embedded_guidance_scale: Optional[float] = None,
    compute_ssim: bool = False,
    compute_lpips: bool = False,
    log_level: str = "INFO",
) -> dict:
    """Compare inference between fine-tuned model and merged/base+adapter model.
    
    Args:
        base: Base or merged model ID/path
        ft: Fine-tuned model ID/path
        adapter: LoRA adapter path (or NONE for merged model)
        output_dir: Output directory for videos
        prompt: Generation prompt
        seed: Random seed
        height: Video height
        width: Video width
        num_frames: Number of frames
        num_inference_steps: Inference steps
        guidance_scale: CFG scale
        flow_shift: Flow shift
        embedded_guidance_scale: Embedded guidance scale
        compute_ssim: Compute SSIM metric
        compute_lpips: Compute LPIPS metric
        log_level: Logging level
        
    Returns:
        Dictionary with metric results
    """
    configure_logging(log_level)

    os.makedirs(output_dir, exist_ok=True)
    try:
        adapter_path = _validate_adapter(adapter)
    except Exception as exc:
        logger.error("Adapter validation failed: %s", exc)
        sys.exit(2)

    # 1) generate with fine-tuned model (reference)
    logger.info("Generating reference (fine-tuned): %s", ft)
    try:
        ft_video = generate_with_model(
            model_path=ft,
            output_dir=output_dir,
            output_name="fine_tuned",
            prompt=prompt,
            seed=seed,
            lora_path=None,
            height=height,
            width=width,
            num_frames=num_frames,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            flow_shift=flow_shift,
            embedded_guidance_scale=embedded_guidance_scale,
        )
        logger.info("Reference video saved: %s", ft_video)
    except Exception as exc:
        logger.error("Reference generation failed: %s", exc)
        sys.exit(3)

    # 2) generate with merged model OR base + adapter
    use_merged = adapter_path is None
    mode = "merged model" if use_merged else "base+adapter"
    logger.info("Generating target (%s): %s", mode, base)
    try:
        target_video = generate_with_model(
            model_path=base,
            output_dir=output_dir,
            output_name="merged_model" if use_merged else "base_plus_lora",
            prompt=prompt,
            seed=seed,
            lora_path=None if use_merged else adapter_path,
            height=height,
            width=width,
            num_frames=num_frames,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            flow_shift=flow_shift,
            embedded_guidance_scale=embedded_guidance_scale,
        )
        logger.info("Target video saved: %s", target_video)
    except Exception as exc:
        logger.error("Target generation failed: %s", exc)
        sys.exit(4)

    # 3) compute metrics
    results = {}
    if compute_ssim or compute_lpips:
        results = compute_metrics(output_dir, ft_video, target_video, num_inference_steps, prompt, compute_ssim, compute_lpips)
        if results.get("mean_ssim") is not None:
            logger.info("Mean SSIM: %.4f", results["mean_ssim"])
        if results.get("mean_lpips") is not None:
            logger.info("Mean LPIPS: %.4f", results["mean_lpips"])

    logger.info("Comparison complete. Videos in: %s", output_dir)
    return results


def main() -> None:
    """CLI wrapper for compare_inference."""
    args = parse_args()
    compare_inference(
        base=args.base,
        ft=args.ft,
        adapter=args.adapter,
        output_dir=args.output_dir,
        prompt=args.prompt,
        seed=args.seed,
        height=args.height,
        width=args.width,
        num_frames=args.num_frames,
        num_inference_steps=args.num_inference_steps,
        guidance_scale=args.guidance_scale,
        flow_shift=args.flow_shift,
        embedded_guidance_scale=args.embedded_guidance_scale,
        compute_ssim=args.compute_ssim,
        compute_lpips=args.compute_lpips,
        log_level=args.log_level,
    )


if __name__ == "__main__":
    main()