#!/usr/bin/env python3
"""Convert GilgameshYX InverseRenderer-1024 into BiliSakura IntrisicWeather-diffusers layout."""

from __future__ import annotations

import json
import shutil
import sys
from pathlib import Path

from diffusers.models.transformers import SD3Transformer2DModel

COLLECTION_ROOT = Path(__file__).resolve().parent
INTRINSIC_REPO = Path("/data/projects/IntrinsicWeather-diffusers")
sys.path.insert(0, str(INTRINSIC_REPO / "src"))
sys.path.insert(0, str(INTRINSIC_REPO))

from intrinsic_weather.models.transformers.transformer_intrinsic_weather import (  # noqa: E402
    IntrinsicWeatherSD3Transformer2DModel,
)
from scripts._conversion_utils import (  # noqa: E402
    ROOT as REPO_ROOT,
    expand_sd3_input_projection,
    merge_sharded_state_dict,
    save_imaa_bundle,
    write_scheduler_config,
)

SD3_PATH = Path(
    "/data/projects/Visual-Generative-Foundation-Model-Collection/models/stabilityai/stable-diffusion-3-medium-diffusers"
)
SD35_TRANSFORMER_REPO = "stabilityai/stable-diffusion-3.5-medium"
CKPT_PATH = Path(
    "/data/projects/Visual-Generative-Foundation-Model-Collection/models/GilgameshYX/InverseRenderer-1024"
)
OUTPUT_ROOT = COLLECTION_ROOT
TRANSFORMER_VARIANT = "inverse-1024"
SHARED_COMPONENTS = (
    "text_encoder",
    "text_encoder_2",
    "text_encoder_3",
    "tokenizer",
    "tokenizer_2",
    "tokenizer_3",
    "vae",
    "scheduler",
)


def copy_sd3_shared_components(sd3_path: Path, output_path: Path) -> None:
    for name in SHARED_COMPONENTS:
        src = sd3_path / name
        dst = output_path / name
        if dst.exists():
            print(f"Skipping existing shared component: {dst}")
            continue
        print(f"Copying {name} ...")
        shutil.copytree(src, dst)


def main() -> None:
    transformer_dir = OUTPUT_ROOT / "transformer" / TRANSFORMER_VARIANT
    transformer_dir.mkdir(parents=True, exist_ok=True)

    print(f"Ensuring shared SD3 components from {SD3_PATH} ...")
    copy_sd3_shared_components(SD3_PATH, OUTPUT_ROOT)
    write_scheduler_config(OUTPUT_ROOT)

    print("Converting inverse renderer transformer (1024) ...")
    base_transformer = SD3Transformer2DModel.from_config(
        SD3Transformer2DModel.load_config(SD35_TRANSFORMER_REPO, subfolder="transformer")
    )
    base_transformer = expand_sd3_input_projection(base_transformer, in_channels=32)
    custom_blocks = IntrinsicWeatherSD3Transformer2DModel.from_config(base_transformer.config)
    custom_blocks.load_state_dict(
        merge_sharded_state_dict(
            [
                CKPT_PATH / "pytorch_model-00001-of-00002.bin",
                CKPT_PATH / "pytorch_model-00002-of-00002.bin",
            ]
        ),
        strict=True,
    )
    custom_blocks.save_pretrained(transformer_dir.as_posix(), safe_serialization=True)
    shutil.copy2(
        REPO_ROOT / "src" / "intrinsic_weather" / "models" / "transformers" / "transformer_intrinsic_weather.py",
        transformer_dir / "transformer_intrinsic_weather.py",
    )

    print("Saving IMAA weights from InverseRenderer-1024 ...")
    save_imaa_bundle(CKPT_PATH / "imaa.pth", OUTPUT_ROOT, safe_serialization=True)

    conversion_metadata = {
        "task": "inverse_renderer",
        "resolution": 1024,
        "transformer_variant": TRANSFORMER_VARIANT,
        "source_transformer_checkpoints": [
            str((CKPT_PATH / "pytorch_model-00001-of-00002.bin").resolve()),
            str((CKPT_PATH / "pytorch_model-00002-of-00002.bin").resolve()),
        ],
        "source_imaa_checkpoint": str((CKPT_PATH / "imaa.pth").resolve()),
        "sd3_path": str(SD3_PATH.resolve()),
        "sd35_transformer_repo": SD35_TRANSFORMER_REPO,
        "in_channels": 32,
    }
    (OUTPUT_ROOT / "conversion_metadata_inverse_1024.json").write_text(
        json.dumps(conversion_metadata, indent=2) + "\n",
        encoding="utf-8",
    )
    print(f"Saved transformer to: {transformer_dir}")
    print("Load with: load_inverse_pipeline(transformer_subfolder='inverse-1024')")


if __name__ == "__main__":
    main()