# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Lucasgitguy-co
"""Experimental, staged Idol inference for Apple MPS.

Upstream revision: 4702a7f00850e34b6e83a2170cc123c3e9570744.
No weights are bundled. Full-weight inference has NOT been verified by this project.
"""
import gc
import json
import os
import shutil
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from .engine import device_for


def patch_attention(layers, query_chunk=128):
    """Bound attention working memory by splitting only the query dimension.

    Each query still sees ALL keys/values. Does not split or change softmax's key axis.
    Disable CUDA-specific optimized sampling and avoid complex RoPE tensors on MPS.
    """
    if query_chunk < 1:
        raise ValueError("query_chunk must be positive")
    from torch.nn import functional as F

    def attention(*args):
        q, k, v = args[0].unbind(dim=2) if len(args) == 1 else args
        q, k, v = (t.permute(0, 2, 1, 3) for t in (q, k, v))
        chunks = [F.scaled_dot_product_attention(q[:, :, start:start+query_chunk], k, v)
                  for start in range(0, q.shape[2], query_chunk)]
        return torch.cat(chunks, dim=2).permute(0, 2, 1, 3)

    layers.scaled_dot_product_attention = attention
    layers.RePo3DRotaryEmbedding.real_valued = True
    return attention


def release(device):
    gc.collect()
    if device.type == "mps":
        torch.mps.synchronize()
        torch.mps.empty_cache()


@torch.no_grad()
def run_idol(args):
    if args.steps < 1 or args.gaussians < 32768 or args.gaussians > 1048576 or args.gaussians % 32:
        raise ValueError("Steps must be positive; Idol requires 32768..1048576 Gaussians, divisible by 32")
    device = device_for(args.device)
    cache = Path(args.cache).resolve()
    cache.mkdir(parents=True, exist_ok=True)
    os.environ["HF_HOME"] = str(cache)
    os.environ["HF_HUB_OFFLINE"] = "0" if args.allow_download else "1"
    os.environ["ATTN_BACKEND"] = "sdpa"
    if args.allow_download and shutil.disk_usage(cache).free < 16 * 1024**3:
        raise ValueError("Less than 16 GiB disk space available. This conservative download guard leaves room for teacher weights, encoders and temporary files. Free space, or use an existing complete cache without --allow-download.")
    try:
        from huggingface_hub import hf_hub_download, snapshot_download
        from safetensors.torch import load_file
        from tensordict import TensorDict
        from idol.conditioning import ImageEncoder, require_encoder_access
        from idol.models import BiRefNet, GaussianFixedlenDecoder, LatentSeqMMFlowModel, OctreeProbabilityFixedlenDecoder, layers
        from idol.preprocess import CROP_MARGIN, preprocess_image, has_real_alpha
        from idol.sampler import sample_flow
        from idol.camera import decode_camera
    except ImportError as exc:
        raise RuntimeError("Install the optional teacher dependencies: python -m pip install '.[teacher]'") from exc
    out = Path(args.out)
    if out.exists():
        raise ValueError("Teacher output directory already exists; choose a new --out")
    # Local-only by default. Hugging Face authentication is never written to project files.
    config_path = hf_hub_download(args.weights, "config.json", local_files_only=not args.allow_download)
    config = json.loads(Path(config_path).read_text())
    encoders = config["image_encoders"]
    require_encoder_access([encoders["dinov3"], encoders["flux2_vae"]])
    weights = Path(snapshot_download(args.weights, allow_patterns=["config.json", "*.safetensors"],
                                    local_files_only=not args.allow_download))
    patch_attention(layers, args.query_chunk)
    torch.manual_seed(args.seed)

    def load(model, name):
        state = load_file(str(weights / f"{name}.safetensors"), device="cpu")
        model.load_state_dict(state, strict=True)
        del state
        return model.eval().requires_grad_(False).to(device)

    print("Experimental Idol MPS adapter: full-weight compatibility and quality remain unverified.", flush=True)
    with Image.open(args.image) as opened:
        image = opened.copy()
    background = None
    if not has_real_alpha(image):
        print("Stage 1: background removal", flush=True)
        background = load(BiRefNet(), "birefnet")
    prepared = preprocess_image(image, background, config["image_size"], CROP_MARGIN[args.mode])
    del background
    release(device)
    pixels = torch.from_numpy(np.array(prepared)).permute(2, 0, 1).float()[None].div(255).to(device)
    print("Stage 2: image conditioning", flush=True)
    encoder = ImageEncoder(encoders["dinov3"], encoders["flux2_vae"], device="cpu")
    # Both encoders initially reside on CPU; only the active encoder is moved to MPS.
    encoder.dinov3.to(device)
    feature1 = encoder._dinov3_tokens(pixels)
    del encoder.dinov3
    release(device)
    encoder.flux2_vae.to(device)
    feature2 = encoder._flux2_tokens(pixels)
    del encoder, pixels
    release(device)
    cond = TensorDict({"feature1": feature1, "feature2": feature2}, batch_size=1, device=device)
    print("Stage 3: latent sampling", flush=True)
    model = load(LatentSeqMMFlowModel(**config["denoiser"]["args"]), "denoiser")
    noise = TensorDict({"latent": torch.randn(1, model.q_token_length, model.in_channels, device=device),
                        "camera": torch.randn(1, 1, model.cam_channels, device=device)}, batch_size=1, device=device)
    result = sample_flow(model, noise, cond=cond, neg_cond=cond.apply(torch.zeros_like), steps=args.steps,
                         rescale_t=3., cfg_strength=3., verbose=True, batch_cfg=False)
    latent, camera = result["latent"], decode_camera(result["camera"][0, 0])
    del model, noise, cond, feature1, feature2, result
    release(device)
    print("Stage 4: octree decoding", flush=True)
    model = load(OctreeProbabilityFixedlenDecoder(**config["decoder"]["args"]), "decoder")
    per_anchor = config["decoder_gs"]["args"]["representation_config"]["num_gaussians"]
    anchors = OctreeProbabilityFixedlenDecoder.sample(model, latent, num_points=args.gaussians // per_anchor,
                                                     level=config["max_voxel_level"])
    del model
    release(device)
    print("Stage 5: Gaussian decoding", flush=True)
    model = load(GaussianFixedlenDecoder(**config["decoder_gs"]["args"]), "decoder_gs")
    gaussian = model.to_representation(anchors, model(anchors, latent))[0]
    out.mkdir(parents=True)
    gaussian.save_ply(str(out / "teacher.ply"))
    prepared.save(out / "input.png")
    (out / "metadata.json").write_text(json.dumps({"teacher": args.weights, "steps": args.steps,
        "gaussians": args.gaussians, "seed": args.seed, "device": str(device), "camera": camera,
        "sdk_revision": "4702a7f00850e34b6e83a2170cc123c3e9570744", "weights_snapshot": weights.name,
        "coordinate_frame": "Idol exported PLY (y-up)", "adapter": "experimental staged MPS"}, indent=2))
    print(f"Teacher output saved to {out}", flush=True)