#!/usr/bin/env python3
"""Faz 2: hibrit kopru + gercek TRELLIS/Hunyuan safetensors uzerinde LoRA."""
from __future__ import annotations

from pathlib import Path
from typing import Any

import torch
import torch.nn as nn

from meshai_train.base_weights import (
    BaseLoRATower,
    _load_safetensors,
    pick_lora_targets,
)
from meshai_train.models import GEOM_IN_DIM, TEXTURE_LATENT_DIM, MeshAIHybridTrainBundle

FAZ2_VERSION = "v5.0-faz2-lora-base"


class MeshAIFaz2Bundle(nn.Module):
    """
    Faz1 hibrit (geometry/texture/bridge) + TRELLIS/Hunyuan LoRA kuleleri.
    Base agirliklar frozen; sadece LoRA + proj + hibrit egitilir.
    """

    def __init__(
        self,
        *,
        trellis_targets: list[tuple[str, torch.Tensor]],
        hunyuan_targets: list[tuple[str, torch.Tensor]],
        lora_rank: int = 8,
    ) -> None:
        super().__init__()
        self.hybrid = MeshAIHybridTrainBundle()
        self.trellis_tower = BaseLoRATower(
            trellis_targets,
            input_dim=GEOM_IN_DIM,
            rank=lora_rank,
            out_dim=TEXTURE_LATENT_DIM,
        )
        # Hunyuan: view embedding -> tower
        self.view_pool = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1),
            nn.GELU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.view_proj = nn.Linear(32, GEOM_IN_DIM)
        self.hunyuan_tower = BaseLoRATower(
            hunyuan_targets,
            input_dim=GEOM_IN_DIM,
            rank=lora_rank,
            out_dim=TEXTURE_LATENT_DIM,
        )

    def forward(self, geom_in: torch.Tensor, views: torch.Tensor) -> dict[str, torch.Tensor]:
        hy = self.hybrid(geom_in, views)
        trellis_feat = self.trellis_tower(geom_in)

        b, v, c, h, w = views.shape
        pooled = self.view_pool(views.reshape(b * v, c, h, w)).reshape(b, v, -1).mean(dim=1)
        view_cond = self.view_proj(pooled)
        hunyuan_feat = self.hunyuan_tower(view_cond)

        return {
            **hy,
            "trellis_feat": trellis_feat,
            "hunyuan_feat": hunyuan_feat,
        }


def faz2_loss(
    out: dict[str, torch.Tensor],
    batch: dict[str, torch.Tensor],
) -> tuple[torch.Tensor, dict[str, float]]:
    voxel_loss = nn.functional.mse_loss(out["voxel_pred"], batch["voxel_tgt"])
    bridge_loss = nn.functional.mse_loss(out["bridge_out"], out["tex_latent"].detach())
    # Gercek base kosullama: LoRA ciktilari hibrit latente hizalansin
    trellis_align = nn.functional.mse_loss(out["trellis_feat"], out["bridge_out"].detach())
    hunyuan_align = nn.functional.mse_loss(out["hunyuan_feat"], out["tex_latent"].detach())
    tex_reg = out["tex_latent"].pow(2).mean() * 1e-4
    total = voxel_loss + bridge_loss + 0.5 * trellis_align + 0.5 * hunyuan_align + tex_reg
    return total, {
        "voxel": float(voxel_loss.item()),
        "bridge": float(bridge_loss.item()),
        "trellis_align": float(trellis_align.item()),
        "hunyuan_align": float(hunyuan_align.item()),
        "tex_reg": float(tex_reg.item()),
    }


def build_faz2_from_weight_files(
    weight_paths: dict[str, Path],
    *,
    lora_rank: int = 8,
    trellis_mats: int = 4,
    hunyuan_mats: int = 4,
    log_fn: Any = print,
) -> MeshAIFaz2Bundle:
    trellis_targets: list[tuple[str, torch.Tensor]] = []
    hunyuan_targets: list[tuple[str, torch.Tensor]] = []

    for rel, path in weight_paths.items():
        state = _load_safetensors(path)
        picks = pick_lora_targets(state, max_matrices=trellis_mats if "TRELLIS" in rel else hunyuan_mats)
        log_fn(f"[faz2] {rel}: {len(state)} tensor, LoRA aday={len(picks)}")
        if "TRELLIS" in rel:
            trellis_targets.extend(picks)
        else:
            hunyuan_targets.extend(picks)

    if not trellis_targets:
        raise RuntimeError("TRELLIS LoRA hedefi bulunamadi")
    if not hunyuan_targets:
        raise RuntimeError("Hunyuan LoRA hedefi bulunamadi")

    # En buyuk N
    trellis_targets = sorted(trellis_targets, key=lambda x: x[1].numel(), reverse=True)[:trellis_mats]
    hunyuan_targets = sorted(hunyuan_targets, key=lambda x: x[1].numel(), reverse=True)[:hunyuan_mats]
    log_fn(
        f"[faz2] secilen TRELLIS mats={len(trellis_targets)} "
        f"Hunyuan mats={len(hunyuan_targets)} rank={lora_rank}"
    )
    return MeshAIFaz2Bundle(
        trellis_targets=trellis_targets,
        hunyuan_targets=hunyuan_targets,
        lora_rank=lora_rank,
    )


def save_faz2_checkpoint(
    path: Path,
    *,
    epoch: int,
    global_step: int,
    model: MeshAIFaz2Bundle,
    extra: dict[str, Any] | None = None,
) -> None:
    payload = {
        "version": FAZ2_VERSION,
        "epoch": epoch,
        "global_step": global_step,
        "model": model.state_dict(),
        "extra": extra or {},
    }
    torch.save(payload, path)


def load_faz2_checkpoint(path: Path, model: MeshAIFaz2Bundle, device: str) -> int:
    state = torch.load(path, map_location=device, weights_only=False)
    if state.get("version") != FAZ2_VERSION:
        # Faz1 hibrit ckpt ise sadece hybrid yukle
        if "geometry" in state and "texture" in state and "bridge" in state:
            model.hybrid.geometry.load_state_dict(state["geometry"])
            model.hybrid.texture.load_state_dict(state["texture"])
            model.hybrid.bridge.load_state_dict(state["bridge"])
            return int(state.get("global_step", 0))
        return 0
    model.load_state_dict(state["model"], strict=False)
    return int(state.get("global_step", 0))