import os

os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import random
import tempfile

import numpy as np
import imageio.v3 as iio
import spaces
import torch
import gradio as gr
from PIL import Image, ImageOps
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from diffusers import LTX2InContextPipeline
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
from diffusers.utils import load_video, encode_video

# --- Config -----------------------------------------------------------------
# Water-simulation IC-LoRA — distilled recipe (8 sigmas, CFG off), strength sweet-spot ~1.2.
BASE_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
LORA_REPO = "ltx-community/LTX-2.3-loras"
LORA_FILE = "ltx-2.3-22b-ic-lora-water-simulation-0.9.safetensors"
LORA_SCALE = 1.2
FPS = 24
NUM_STEPS = len(DISTILLED_SIGMA_VALUES)
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")

RES_PRESETS = {"Fast (768×448)": (768, 448), "Quality (960×544)": (960, 544)}
FRAME_CHOICES = [49, 73, 97, 121]

pipe = LTX2InContextPipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe.vae.enable_tiling()
_lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=HF_TOKEN)
pipe.load_lora_weights(load_file(_lora_path), adapter_name="water")
pipe.set_adapters("water", LORA_SCALE)


def _src_fps(path, default=FPS):
    try:
        return float(iio.immeta(path, plugin="pyav").get("fps", default)) or default
    except Exception:
        return default


def _load_frames(path, num_frames, width, height):
    frames = load_video(path)
    if not frames:
        return []
    fps = _src_fps(path)
    out = []
    for i in range(num_frames):
        idx = min(int(round(i / FPS * fps)), len(frames) - 1)
        out.append(ImageOps.fit(frames[idx].convert("RGB"), (width, height), Image.LANCZOS))
    return out


def _pick_resolution(first_frame, preset):
    w, h = RES_PRESETS[preset]
    if first_frame.height > first_frame.width:
        w, h = h, w
    return w, h


def _build_prompt(prompt):
    desc = prompt.strip() or "a flowing stream of clear water"
    return (
        "Reference shows the dry scene. Edited shows the same scene with water added. "
        f"ADD WATER {desc}. "
        "Subject identity, clothing, framing, and background geometry are identical to the reference; "
        "only water-related elements differ between reference and edited."
    )


def _export(video_np, audio, path):
    kw = {}
    if audio is not None:
        kw = dict(audio=audio[0].float().cpu(), audio_sample_rate=pipe.vocoder.config.output_sampling_rate)
    encode_video(video_np, fps=FPS, output_path=path, **kw)


def _duration(*args, **kwargs):
    preset = next((a for a in args if isinstance(a, str) and a in RES_PRESETS), "Fast")
    num_frames = next((a for a in args if isinstance(a, int) and a in FRAME_CHOICES), 73)
    per_frame = 1.6 if "Quality" in str(preset) else 1.0
    return int(70 + int(num_frames) * per_frame)


@spaces.GPU(duration=_duration)
def add_water(video, prompt, strength, preset, num_frames, seed, randomize,
              progress=gr.Progress(track_tqdm=True)):
    if video is None:
        raise gr.Error("Please upload a 'dry' video to add water to.")
    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    num_frames = int(num_frames)

    probe = load_video(video)
    if not probe:
        raise gr.Error("Could not read any frames from that video.")
    width, height = _pick_resolution(probe[0], preset)
    ref = _load_frames(video, num_frames, width, height)
    pipe.set_adapters("water", float(strength))
    full_prompt = _build_prompt(prompt)

    def _cb(p, i, t, kw):
        progress((i + 1) / NUM_STEPS, desc=f"Adding water — step {i + 1}/{NUM_STEPS}")
        return {}

    video_out, audio_out = pipe(
        prompt=full_prompt, negative_prompt="",
        reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
        reference_downscale_factor=1,
        width=width, height=height, num_frames=num_frames, frame_rate=FPS,
        num_inference_steps=NUM_STEPS, sigmas=DISTILLED_SIGMA_VALUES,
        guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
        generator=torch.Generator(device="cuda").manual_seed(seed),
        output_type="np", return_dict=False, callback_on_step_end=_cb,
    )
    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    _export(video_out[0], audio_out, out_path)
    return out_path, seed


with gr.Blocks(title="LTX-2.3 Water Simulation") as demo:
    gr.Markdown(
        "# 🌊 LTX-2.3 Water Simulation\n"
        "Add believable, naturally-moving water to a dry clip — rivers, surf, rain, waterfalls, floods, "
        "splashes — that interacts with the moving scene, while subject, clothing, framing and camera stay "
        "exactly as shot. Describe the water (and any sounds) in one prompt; the trigger `ADD WATER` is added for you. "
        "IC-LoRA: [`ltx-community/LTX-2.3-loras`](https://huggingface.co/ltx-community/LTX-2.3-loras) · base: distilled LTX-2.3."
    )
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Dry input video")
            prompt = gr.Textbox(
                label="Describe the water — type, motion, how it interacts, plus any sounds", lines=3,
                placeholder="a clear shallow stream braiding around their legs with white foam crests and glistening wet ground; rushing water, gentle splashing",
            )
            with gr.Accordion("Settings", open=False):
                strength = gr.Slider(1.0, 1.6, value=1.2, step=0.05,
                                     label="Water strength (1.2–1.3 natural · 1.35+ hard surface→sea · ≥1.5 max drama)")
                preset = gr.Dropdown(list(RES_PRESETS), value="Fast (768×448)", label="Resolution")
                num_frames = gr.Dropdown(FRAME_CHOICES, value=73, label="Frames (24fps)")
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            run = gr.Button("Add water", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Result with water")
            used_seed = gr.Number(label="Seed used", interactive=False)

    run.click(add_water, inputs=[video_in, prompt, strength, preset, num_frames, seed, randomize],
              outputs=[video_out, used_seed])

    gr.Examples(
        examples=[
            ["examples/man_dancing_dry.mp4",
             "a clear shallow stream rushing and braiding around their legs with white foam crests and glistening wet floor, splashing with each step; rushing water and rhythmic splashes",
             1.3, "Fast (768×448)", 73, 42, False],
            ["examples/landscape_dry.mp4",
             "a wide river flooding across the valley with rippling reflections and drifting foam, mist rising off the surface; flowing water and a distant waterfall",
             1.25, "Fast (768×448)", 73, 42, False],
        ],
        inputs=[video_in, prompt, strength, preset, num_frames, seed, randomize],
        outputs=[video_out, used_seed], fn=add_water, cache_examples=True, cache_mode="lazy",
    )

if __name__ == "__main__":
    demo.launch(show_error=True)