"""Full inference for HunyuanPaintPBR with CFG + multi-view.

Reproduces the PyTorch pipeline's denoising loop:
- 6-view generation (albedo + metallic-roughness)
- Classifier-free guidance (triple-batch: uncond / ref / full)
- Learned material text embeddings
- DINO feature conditioning
- View-dependent guidance scaling
"""

from typing import Dict, List, Optional

import mlx.core as mx
import mlx.nn as nn
import numpy as np
from PIL import Image

from .dino_mlx import preprocess_for_dino
from .load_model import HunyuanPaintModelMLX, extract_reference_features


def _compute_dino_features_hf(model, ref_u8: np.ndarray) -> mx.array:
    """Run HF DINOv2-giant once on the reference image.

    Uses fp16 on CPU (1.2GB vs 4.4GB fp32) to avoid OOM with the paint
    UNet + dual UNet + VAE + super-res all resident. Frees the model
    after feature extraction.

    Falls back to pure MLX DINO if HF transformers isn't available.
    """
    try:
        import torch
        import gc as _gc
        from transformers import AutoImageProcessor, AutoModel
        from PIL import Image as PILImage
    except Exception:
        from .dino_mlx import preprocess_for_dino
        return model.dino(preprocess_for_dino(ref_u8))

    proc = AutoImageProcessor.from_pretrained("facebook/dinov2-giant")
    pt_dino = AutoModel.from_pretrained(
        "facebook/dinov2-giant", torch_dtype=torch.float16,
    )
    for p in pt_dino.parameters():
        p.requires_grad_(False)
    pt_in = proc(images=PILImage.fromarray(ref_u8), return_tensors="pt")
    with torch.no_grad():
        pt_feat = pt_dino(pt_in.pixel_values.half())[0]
    out = mx.array(pt_feat.float().numpy())
    del proc, pt_dino, pt_in, pt_feat
    _gc.collect()
    return out


def _cam_mapping(azim: float) -> float:
    """View-dependent guidance scale (per the original pipeline)."""
    azim = azim % 360
    if 0 <= azim < 90:
        return azim / 90.0 + 1.0
    elif 90 <= azim < 330:
        return 2.0
    else:
        return -azim / 90.0 + 5.0


def generate_multiview_pbr(
    model: HunyuanPaintModelMLX,
    normal_maps: List[np.ndarray],
    position_maps: List[np.ndarray],
    reference_image: np.ndarray,
    camera_azims: Optional[List[float]] = None,
    num_inference_steps: int = 15,
    guidance_scale: float = 3.0,
    view_size: int = 512,
    seed: int = 0,
) -> Dict[str, List[np.ndarray]]:
    """Generate multiview PBR textures.

    Args:
        model: Loaded HunyuanPaintModelMLX.
        normal_maps: List of N (H, W, 3) float32 normal renders in [0, 1].
        position_maps: List of N (H, W, 3) float32 position renders in [0, 1].
        reference_image: (H, W, 3) uint8 or float32 reference image.
        camera_azims: Azimuth angles for each view (for view-dependent guidance).
        num_inference_steps: Number of denoising steps.
        guidance_scale: CFG scale.
        view_size: Generation resolution.
        seed: Random seed.

    Returns:
        dict with 'albedo' and 'mr' keys, each a list of N (H, W, 3) uint8 arrays.
    """
    mx.random.seed(seed)

    n_views = len(normal_maps)
    n_pbr = 2  # albedo + mr
    latent_size = view_size // 8
    scaling_factor = 0.18215

    if camera_azims is None:
        camera_azims = [0.0] * n_views

    # ------------------------------------------------------------------
    # 1. Encode conditions
    # ------------------------------------------------------------------
    print("  Encoding conditions...")

    def encode_images(images):
        """Encode list of images to latent space."""
        arrays = []
        for img in images:
            if img.dtype == np.uint8:
                img = img.astype(np.float32) / 255.0
            from PIL import Image as PILImage
            pil = PILImage.fromarray((img * 255).astype(np.uint8)).resize(
                (view_size, view_size)
            )
            arr = np.array(pil).astype(np.float32) / 255.0
            arrays.append(arr)
        batch = mx.array(np.stack(arrays))  # (N, H, W, 3)
        # vae.encode already multiplies by self.scaling_factor internally;
        # do NOT re-multiply (was double-scaling latents to ~0.033 of the
        # expected magnitude, making ref_latents and z_normal/z_pos
        # essentially invisible to the UNet → washed-out diffusion output).
        return model.vae.encode(batch * 2 - 1)

    z_normal = encode_images(normal_maps)   # (N, h, w, 4)
    z_pos = encode_images(position_maps)    # (N, h, w, 4)
    mx.synchronize()

    # Multi-resolution voxel indices for 3D RoPE in multiview attention.
    # Built from the raw position maps (NOT the VAE-encoded latents).
    from .unet.voxel_indices import calc_multires_voxel_idxs
    pos_resized = []
    for pm in position_maps:
        if pm.dtype == np.uint8:
            pm = pm.astype(np.float32) / 255.0
        from PIL import Image as PILImage
        pil = PILImage.fromarray((pm * 255).astype(np.uint8)).resize(
            (view_size, view_size)
        )
        pos_resized.append(np.array(pil).astype(np.float32) / 255.0)
    # (N, H, W, 3) -> (1, N, 3, H, W)
    pos_np = np.stack(pos_resized).transpose(0, 3, 1, 2)[None]
    voxel_indices = calc_multires_voxel_idxs(pos_np)

    # ------------------------------------------------------------------
    # 2. DINO features
    # ------------------------------------------------------------------
    print("  Extracting DINO features...")
    if reference_image.dtype != np.uint8:
        ref_u8 = (reference_image * 255).astype(np.uint8)
    else:
        ref_u8 = reference_image
    # Pure MLX DINO at native 518 (1370 tokens). Pool spatial grid from
    # 37x37 down to 16x16 so the final token count (257) matches what
    # HF DINOv2 produces at 224 input — the distribution attn_dino's
    # to_k / to_v were trained against. HF DINO bridge OOMs 32 GB in
    # both fp32 and fp16.
    dino_in = preprocess_for_dino(ref_u8)
    dino_feat = model.dino(dino_in)  # (1, 1370, 1536)
    # Separate CLS + patch grid, avg-pool 37x37 -> 16x16
    import math as _math
    N_full = dino_feat.shape[1] - 1
    g_old = int(round(_math.sqrt(N_full)))
    if g_old * g_old == N_full and g_old > 16:
        cls = dino_feat[:, :1, :]
        patch = dino_feat[:, 1:, :].reshape(1, g_old, g_old, -1)
        # Downsample to 16x16 via adaptive average pool (bilinear in MLX)
        up = nn.Upsample(scale_factor=(16.0 / g_old, 16.0 / g_old),
                         mode="linear")
        patch_d = up(patch).reshape(1, 16 * 16, -1)
        dino_feat = mx.concatenate([cls, patch_d], axis=1)
    dino_proj = model.image_proj(dino_feat)
    mx.synchronize()

    # ------------------------------------------------------------------
    # 3. Text embeddings for CFG
    # ------------------------------------------------------------------
    # Per-material learned tokens
    text_albedo = model.learned_text_clip["albedo"]  # (77, 1024)
    text_mr = model.learned_text_clip["mr"]
    # PT (pipeline.py:275) sets ``negative_prompt_embeds = prompt_embeds``
    # when use_learned_text_clip is True — the uncond batch reuses the
    # SAME learned material tokens, not zeros. CFG signal then comes
    # from ref_features/DINO presence, not from text. Using zeros here
    # was pushing the uncond prediction off-distribution and weakening
    # the guidance.

    # ------------------------------------------------------------------
    # 2b. Reference features — capture norm_hidden_states from each
    #     transformer block by running a single forward pass on the
    #     reference image latent.  Matches PyTorch mode "w".
    # ------------------------------------------------------------------
    print("  Extracting reference features...")
    ref_latent = encode_images([reference_image])  # (1, h, w, 4)
    text_ref = model.learned_text_clip["ref"]  # PT uses learned_text_clip_ref
    ref_text_batch = text_ref[None]  # (1, 77, 1024)
    # Use the dual-stream reference UNet when available — its weights are
    # different from the main UNet and are specifically trained to extract
    # ref-attention features. Falling back to the main UNet (as we did) gives
    # noticeably worse reference grounding.
    ref_unet = getattr(model, "unet_dual", None) or model.unet
    ref_features = extract_reference_features(
        ref_unet, ref_latent, ref_text_batch
    )
    mx.synchronize()

    # Build per-sample text: (N_pbr * N_views, 77, 1024)
    # Order: [albedo_v0, albedo_v1, ..., mr_v0, mr_v1, ...]
    text_full = mx.concatenate([
        mx.broadcast_to(text_albedo[None], (n_views, 77, 1024)),
        mx.broadcast_to(text_mr[None], (n_views, 77, 1024)),
    ], axis=0)  # (N_pbr*N_views, 77, 1024) = (12, 77, 1024)

    # Negative and positive text batches are identical (PT pipeline.py:275).
    text_neg_batch = text_full

    # CFG triple: [negative, ref(=full), full]
    text_cfg = mx.concatenate([text_neg_batch, text_full, text_full], axis=0)
    # (36, 77, 1024)

    # DINO features for CFG: replicate for all 36 samples
    # But set to zero for negative (uncond) batch
    dino_for_cfg = mx.broadcast_to(
        dino_proj, (n_pbr * n_views, dino_proj.shape[1], dino_proj.shape[2])
    )
    dino_zeros = mx.zeros_like(dino_for_cfg)
    dino_cfg = mx.concatenate([dino_zeros, dino_for_cfg, dino_for_cfg], axis=0)

    # ------------------------------------------------------------------
    # 4. View-dependent guidance scale
    # ------------------------------------------------------------------
    view_scales = [_cam_mapping(a) for a in camera_azims]
    # Repeat for materials: [albedo_scales, mr_scales]
    view_scale_flat = np.array(view_scales * n_pbr, dtype=np.float32)
    view_scale = mx.array(view_scale_flat)[:, None, None, None]  # (12, 1, 1, 1)

    # ------------------------------------------------------------------
    # 5. Prepare condition latents for concat
    # ------------------------------------------------------------------
    # Replicate conditions for N_pbr materials
    z_normal_rep = mx.concatenate([z_normal] * n_pbr, axis=0)  # (12, h, w, 4)
    z_pos_rep = mx.concatenate([z_pos] * n_pbr, axis=0)        # (12, h, w, 4)

    # ------------------------------------------------------------------
    # 6. Denoising loop
    # ------------------------------------------------------------------
    latents = mx.random.normal((n_pbr * n_views, latent_size, latent_size, 4))
    model.scheduler.set_timesteps(num_inference_steps)

    print(f"  Denoising ({num_inference_steps} steps, {n_views} views, CFG={guidance_scale})...")

    for i, t in enumerate(model.scheduler.timesteps):
        t_int = int(t)

        # Process all n_views in a single forward so the multiview
        # attention path inside BasicTransformerBlock fires (it guards on
        # n_views > 1). Combined with the 3D RoPE positional encoding
        # from calc_multires_voxel_idxs and the norm_hs-reuse fix, the
        # cross-view consistency this provides is a big part of matching
        # PT's output.
        chunk_size = n_views
        t_arr = mx.array([t_int])
        noise_guided = mx.zeros_like(latents)

        for v_start in range(0, n_views, chunk_size):
            v_end = min(v_start + chunk_size, n_views)
            n_chunk = v_end - v_start

            # Gather chunk indices for both materials
            # Order: [albedo_v_start..v_end, mr_v_start..v_end]
            # Maps into full latent array ordered as
            # [albedo_0..N-1, mr_0..N-1]
            chunk_idx = list(range(v_start, v_end)) + [
                n_views + j for j in range(v_start, v_end)
            ]
            lat_chunk = latents[mx.array(chunk_idx)]           # (2*n_chunk, h, w, 4)
            zn_chunk = z_normal_rep[mx.array(chunk_idx)]
            zp_chunk = z_pos_rep[mx.array(chunk_idx)]
            unet_in = mx.concatenate([lat_chunk, zn_chunk, zp_chunk], axis=-1)

            n_c = len(chunk_idx)
            text_neg_c = text_neg_batch[:n_c]
            text_full_c = text_full[mx.array(chunk_idx)]
            dino_c = dino_for_cfg[:n_c]

            # 3 CFG passes (sequential to save memory).
            # Mirrors the PyTorch pipeline (hunyuanpaintpbr/pipeline.py:300+,
            # hunyuanpaintpbr/pipeline.py:685):
            #   uncond  -> ref_scale=0, dino=zero, prompt=neg
            #   ref     -> ref_scale=1, dino=zero, prompt=pos
            #   full    -> ref_scale=1, dino=full, prompt=pos
            # Then: noise = uncond + g*v*(ref - uncond) + g*v*(full - ref).
            # Earlier we passed DINO to BOTH ref and full, making
            # (full - ref) ~= 0 and dropping half of the guidance.

            # Slice voxel indices to the current chunk's views
            chunk_view_ids = list(range(v_start, v_end))
            chunk_voxel = {}
            for seq_len, vd in voxel_indices.items():
                vi_full = vd["voxel_indices"]  # (1, n_views * GxG, 3)
                gxg = vi_full.shape[1] // n_views
                # Pick this chunk's views
                vi_per_view = vi_full.reshape(1, n_views, gxg, 3)
                vi_chunk = vi_per_view[:, mx.array(chunk_view_ids), :, :]
                vi_chunk = vi_chunk.reshape(1, len(chunk_view_ids) * gxg, 3)
                chunk_voxel[len(chunk_view_ids) * gxg] = {
                    "voxel_indices": vi_chunk,
                    "voxel_resolution": vd["voxel_resolution"],
                }

            # uncond: no DINO, no ref features, neg text
            pred_uncond = model.unet(
                unet_in, t_arr, text_neg_c,
                n_views=n_chunk, n_pbr=n_pbr,
                position_voxel_indices=chunk_voxel,
            )
            mx.synchronize()

            # ref: ref_features ON, DINO OFF, pos text
            ctx_ref = dict(
                n_views=n_chunk, n_pbr=n_pbr,
                position_voxel_indices=chunk_voxel,
            )
            if ref_features is not None:
                ctx_ref["ref_features"] = ref_features
            pred_ref = model.unet(unet_in, t_arr, text_full_c, **ctx_ref)
            mx.synchronize()

            # full: ref_features ON, DINO ON, pos text
            ctx_full = dict(
                n_views=n_chunk, n_pbr=n_pbr,
                dino_features=dino_c,
                position_voxel_indices=chunk_voxel,
            )
            if ref_features is not None:
                ctx_full["ref_features"] = ref_features
            pred_full = model.unet(unet_in, t_arr, text_full_c, **ctx_full)
            mx.synchronize()

            vs = view_scale[mx.array(chunk_idx)]
            guided = (
                pred_uncond
                + guidance_scale * vs * (pred_ref - pred_uncond)
                + guidance_scale * vs * (pred_full - pred_ref)
            )

            # Scatter chunk results back to correct positions in the
            # full latent array.  chunk_idx maps each chunk sample to
            # its position in the [albedo_0..N-1, mr_0..N-1] layout.
            noise_guided[mx.array(chunk_idx)] = guided

        # Scheduler step
        latents = model.scheduler.step(noise_guided, t_int, latents)
        mx.synchronize()

        if i % 5 == 0 or i == num_inference_steps - 1:
            r = float(latents.max() - latents.min())
            print(f"    step {i}/{num_inference_steps}: t={t_int}, range={r:.1f}")

    # ------------------------------------------------------------------
    # 7. Decode (one image at a time to avoid OOM)
    # ------------------------------------------------------------------
    print("  Decoding...")
    all_images = []
    n_total = n_pbr * n_views
    for idx in range(n_total):
        # latents already at scaling_factor magnitude; vae.decode handles
        # the divide internally — pass z as-is.
        z_i = latents[idx : idx + 1]
        dec_i = model.vae.decode(z_i)
        img_i = mx.clip((dec_i + 1) / 2, 0, 1)
        mx.synchronize()
        all_images.append(np.array(img_i[0]))
        if (idx + 1) % n_views == 0:
            print(f"    decoded {idx + 1}/{n_total}")

    result = {
        "albedo": [
            (all_images[i] * 255).astype(np.uint8)
            for i in range(n_views)
        ],
        "mr": [
            (all_images[n_views + i] * 255).astype(np.uint8)
            for i in range(n_views)
        ],
    }
    return result
