#!/usr/bin/env python3
"""Smoke-test all IntrinsicWeather pipelines on CUDA with bfloat16."""

from __future__ import annotations

import gc
import sys
from pathlib import Path

import torch

REPO = Path(__file__).resolve().parent
sys.path.insert(0, str(REPO))

DTYPE = torch.bfloat16
DEVICE = "cuda"
IMAGE_SIZE = 512
STEPS = 2


def _clear():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def test_inverse() -> None:
    from imaa.imaa import IMAA
    from pipeline_intrinsic_weather_inverse import IntrinsicWeatherInversePipeline
    from safetensors.torch import load_file

    print("[inverse] loading pipeline ...")
    pipe = IntrinsicWeatherInversePipeline.from_pretrained(
        REPO,
        transformer_subfolder="inverse-512",
        device=DEVICE,
        local_files_only=True,
        torch_dtype=DTYPE,
    )
    assert next(pipe.transformer.parameters()).dtype == DTYPE
    assert next(pipe.transformer.parameters()).device.type == "cuda"
    print(f"[inverse] transformer in_channels={pipe.transformer.config.in_channels}")

    imaa = IMAA(dino_model=None, processor=None, num_maps=5, map_embedding_dim=256, common_dim=128).to(DEVICE)
    imaa.load_state_dict(load_file((REPO / "imaa" / "model.safetensors").as_posix()))
    imaa.eval()

    image = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE, device=DEVICE, dtype=DTYPE)
    prompt_embeds, _, pooled_prompt_embeds, _ = pipe.encode_prompt(
        prompt="Albedo (diffuse basecolor)",
        prompt_2=None,
        prompt_3=None,
        do_classifier_free_guidance=False,
    )
    print("[inverse] running 2-step inference ...")
    out = pipe(
        image=image,
        prompt_embeds=prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        guidance_scale=0.0,
        image_guidance_scale=0.0,
        num_inference_steps=STEPS,
        output_type="pt",
        aov=["albedo"],
        map_aware_mask=None,
    )
    print(f"[inverse] output shape={tuple(out.images[0].shape)} dtype={out.images[0].dtype}")
    del pipe, imaa, out
    _clear()


def test_forward() -> None:
    from pipeline_intrinsic_weather_forward import IntrinsicWeatherForwardPipeline

    print("[forward] loading pipeline ...")
    pipe = IntrinsicWeatherForwardPipeline.from_pretrained(
        REPO,
        transformer_subfolder="forward",
        device=DEVICE,
        local_files_only=True,
        torch_dtype=DTYPE,
        load_lora=True,
    )
    assert next(pipe.transformer.parameters()).dtype == DTYPE
    assert next(pipe.transformer.parameters()).device.type == "cuda"
    print(f"[forward] transformer in_channels={pipe.transformer.config.in_channels}")

    aov = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE, device=DEVICE, dtype=DTYPE)
    print("[forward] running 2-step inference ...")
    out = pipe(
        albedo=aov,
        normal=aov,
        roughness=aov,
        metallic=aov,
        irradiance=aov,
        prompt=["A rainy day."],
        guidance_scale=6.0,
        image_guidance_scale=1.5,
        num_inference_steps=STEPS,
        required_aovs=["albedo", "normal", "roughness", "metallic", "irradiance"],
        generator=torch.Generator(device=DEVICE).manual_seed(0),
    )
    print(f"[forward] output size={out.images[0].size}")
    del pipe, out
    _clear()


def test_unified() -> None:
    from pipeline_intrinsic_weather import IntrinsicWeatherPipeline

    print("[unified] loading pipeline ...")
    pipe = IntrinsicWeatherPipeline.from_pretrained(
        REPO,
        inverse_transformer_subfolder="inverse-512",
        forward_transformer_subfolder="forward",
        device=DEVICE,
        local_files_only=True,
        torch_dtype=DTYPE,
        load_lora=True,
        load_imaa=True,
    )
    assert next(pipe.inverse_transformer.parameters()).dtype == DTYPE
    assert next(pipe.forward_transformer.parameters()).dtype == DTYPE
    print("[unified] loaded inverse + forward transformers and IMAA")
    del pipe
    _clear()


def main() -> None:
    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required for this test.")

    print(f"device={torch.cuda.get_device_name(0)} dtype={DTYPE}")
    test_inverse()
    test_forward()
    test_unified()
    print("All pipeline tests passed.")


if __name__ == "__main__":
    main()