"""Dual-encoder модель для symmetric text-guided cross-view geo-localization.

    И дрон, И спутник имеют собственное текстовое описание.
        Слияние происходит СИММЕТРИЧНО на обеих сторонах:

    Drone (картинка + текст → слияние) ──┐
                                          ├─→ cosine similarity → InfoNCE
    Satellite (картинка + текст → слияние) ┘

"""
from __future__ import annotations

import logging
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F

from src.models.dgtrs.model import (
    DGTRSTextEncoder,
    load_dgtrs_text_encoder,
    tokenize_dgtrs,
)
from src.models.stripnet_encoder import StripNetEncoder

LOGGER = logging.getLogger("cvgl.dual_encoder")


# ---------------------------------------------------------------------------
# Слияние картинки и текста (GatedFusion) — каноническая скалярная формула
# ---------------------------------------------------------------------------

class GatedFusion(nn.Module):
    """Каноническое gated-слияние визуальных и текстовых признаков.

        fused = σ(α)·img_proj + (1 − σ(α))·text_proj

    Формула соответствует первоисточнику (GeoRSCLIP GatedFusion) и
    протоколу §3.2: `α` — ЕДИНСТВЕННЫЙ обучаемый скаляр, инициализируется
    так, что σ(α) = init_gate (вес картинки на старте). σ(α) → 1 означает
    "доверяем картинке", σ(α) → 0 — "доверяем тексту".

    Картинка и текст приходят из РАЗНЫХ энкодеров (StripNet 1024-d,
    DGTRS 512-d), поэтому сначала проецируются в общий fused_dim и
    L2-нормируются — чтобы выпуклая комбинация была осмысленной
    (комбинируем единичные векторы, как в общем пространстве CLIP
    первоисточника).

    baseline_mode=True: gate ≡ 1.0 → текст полностью игнорируется
    (нужно для baseline B0 и теста разрушения текста, протокол §8).

    Используется как для дрона, так и для спутника — но как ДВА РАЗНЫХ
    экземпляра с независимыми весами (см. DualEncoder), поскольку визуальные
    домены различаются.

    Args:
        image_dim: Размерность визуальных признаков (вход).
        text_dim:  Размерность текстовых признаков (вход).
        fused_dim: Размерность слитого вектора (выход).
        init_gate: Начальный вес картинки σ(α) ∈ (0, 1).
        baseline_mode: Если True — gate ≡ 1.0 (текст выключен).
    """

    def __init__(
        self,
        image_dim: int,
        text_dim: int,
        fused_dim: int,
        init_gate: float = 0.7,
        baseline_mode: bool = False,
    ) -> None:
        super().__init__()
        if not 0.0 < init_gate < 1.0:
            raise ValueError(f"init_gate must be in (0, 1), got {init_gate}")

        self.image_proj = nn.Linear(image_dim, fused_dim)
        self.text_proj = nn.Linear(text_dim, fused_dim)

        # α в логит-пространстве: σ(α) = init_gate на старте.
        init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
        self.alpha = nn.Parameter(init_alpha)
        self.baseline_mode = baseline_mode

        for m in (self.image_proj, self.text_proj):
            nn.init.trunc_normal_(m.weight, std=0.02)
            nn.init.zeros_(m.bias)

    def forward(
        self,
        image_feat: torch.Tensor,
        text_feat: torch.Tensor | None,
    ) -> torch.Tensor:
        """Слить визуальные и текстовые признаки одной скалярной gate.

        Args:
            image_feat: [B, image_dim] — визуальные признаки (до проекции).
            text_feat:  [B, text_dim]  — текстовые признаки, или None.

        Returns:
            [B, fused_dim] — слитый вектор σ(α)·img + (1 − σ(α))·text.
        """
        img = F.normalize(self.image_proj(image_feat), dim=-1)
        if text_feat is None or self.baseline_mode:
            return img
        txt = F.normalize(self.text_proj(text_feat), dim=-1)
        gate = torch.sigmoid(self.alpha)
        return gate * img + (1.0 - gate) * txt

    @property
    def gate_value(self) -> float:
        """Текущий вес картинки σ(α). 1.0 = текст игнорируется."""
        if self.baseline_mode:
            return 1.0
        return torch.sigmoid(self.alpha).item()


# ---------------------------------------------------------------------------
# DualEncoder: drone (fused) vs satellite (fused) — СИММЕТРИЧНО
# ---------------------------------------------------------------------------

class DualEncoder(nn.Module):
    """Symmetric dual-encoder: слитый drone-эмбеддинг vs слитый satellite-эмбеддинг.

    И дрон, и спутник проходят одинаковую по структуре, но раздельную по
    весам процедуру: картинка + текст → GatedFusion → проекция →
    L2-норма. Полученные эмбеддинги сравниваются через cosine similarity.

    Args:
        text_encoder:              DGTRS-CLIP текстовый энкодер — ОБЩИЙ для
                                     текста дрона и текста спутника (это одна
                                     и та же языковая модель).
        drone_image_encoder:       StripNet энкодер для картинок дрона.
        satellite_image_encoder:   StripNet энкодер для спутниковых снимков
                                     (отдельный экземпляр — домены различны).
        fused_dim:                 Размерность вектора после слияния.
        shared_dim:                Итоговая размерность общего пространства.
        temperature_init:          Начальное значение temperature для InfoNCE.
    """

    def __init__(
        self,
        text_encoder: DGTRSTextEncoder,
        drone_image_encoder: StripNetEncoder,
        satellite_image_encoder: StripNetEncoder,
        fused_dim: int = 512,
        shared_dim: int = 512,
        temperature_init: float = 0.07,
        init_gate: float = 0.7,
        baseline_mode: bool = False,
    ) -> None:
        super().__init__()

        self.text_encoder = text_encoder
        self.drone_image_encoder = drone_image_encoder
        self.satellite_image_encoder = satellite_image_encoder
        self.fused_dim = fused_dim
        self.shared_dim = shared_dim
        self.baseline_mode = baseline_mode

        # --- Слияние картинки + текста — ДВА отдельных модуля ---
        # Раздельные веса: визуальные домены дрона и спутника различны,
        # поэтому оптимальный способ "смешивания" картинки с текстом тоже
        # должен быть индивидуальным для каждой стороны.
        self.drone_fusion = GatedFusion(
            image_dim=drone_image_encoder.out_dim,
            text_dim=text_encoder.embed_dim,
            fused_dim=fused_dim,
            init_gate=init_gate,
            baseline_mode=baseline_mode,
        )
        self.satellite_fusion = GatedFusion(
            image_dim=satellite_image_encoder.out_dim,
            text_dim=text_encoder.embed_dim,
            fused_dim=fused_dim,
            init_gate=init_gate,
            baseline_mode=baseline_mode,
        )

        # --- Проекционные головы в общее пространство сравнения ---
        self.drone_proj = nn.Sequential(
            nn.Linear(fused_dim, shared_dim),
            nn.GELU(),
            nn.Linear(shared_dim, shared_dim),
        )
        self.satellite_proj = nn.Sequential(
            nn.Linear(fused_dim, shared_dim),
            nn.GELU(),
            nn.Linear(shared_dim, shared_dim),
        )
        for proj in (self.drone_proj, self.satellite_proj):
            for m in proj.modules():
                if isinstance(m, nn.Linear):
                    nn.init.trunc_normal_(m.weight, std=0.02)
                    if m.bias is not None:
                        nn.init.zeros_(m.bias)

        # --- Learnable temperature ---
        self.log_temperature = nn.Parameter(
            torch.tensor(temperature_init).log()
        )

    @property
    def temperature(self) -> torch.Tensor:
        return self.log_temperature.exp().clamp(min=0.01, max=0.1)

    def encode_drone(
        self,
        drone_images: torch.Tensor,
        drone_tokens: torch.Tensor,
    ) -> torch.Tensor:
        """Дрон (картинка + текст) → слитый, нормализованный эмбеддинг.

        Args:
            drone_images: [B, 3, H, W] — картинки дрона.
            drone_tokens: [B, 248]     — токенизированный текст описания дрона.
        """
        image_feat = self.drone_image_encoder(drone_images)
        text_feat = self.text_encoder(drone_tokens)
        fused = self.drone_fusion(image_feat, text_feat)
        projected = self.drone_proj(fused)
        return F.normalize(projected, dim=-1)

    def encode_satellite(
        self,
        satellite_images: torch.Tensor,
        satellite_tokens: torch.Tensor,
    ) -> torch.Tensor:
        """Спутник (картинка + текст) → слитый, нормализованный эмбеддинг.

        Args:
            satellite_images: [B, 3, H, W] — спутниковые снимки галереи.
            satellite_tokens: [B, 248]     — токенизированный текст описания
                                              спутникового снимка.
        """
        image_feat = self.satellite_image_encoder(satellite_images)
        text_feat = self.text_encoder(satellite_tokens)
        fused = self.satellite_fusion(image_feat, text_feat)
        projected = self.satellite_proj(fused)
        return F.normalize(projected, dim=-1)

    def forward(
        self,
        drone_images: torch.Tensor,
        drone_tokens: torch.Tensor,
        satellite_images: torch.Tensor,
        satellite_tokens: torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        """Прямой проход: дрон (слитый) vs спутник (слитый).

        Returns:
            dict:
                drone_emb:      [B, shared_dim]
                satellite_emb:  [B, shared_dim]
                logits:         [B, B] — drone_emb @ satellite_emb.T / temperature
                temperature:    скаляр
                drone_gate:     σ(α) дрон-fusion (вес картинки, для логов)
                satellite_gate: σ(α) спутник-fusion (вес картинки, для логов)
        """
        drone_emb = self.encode_drone(drone_images, drone_tokens)
        satellite_emb = self.encode_satellite(satellite_images, satellite_tokens)

        # logits[i][j] = cosine_sim(drone_i, satellite_j) / temperature
        logits = drone_emb @ satellite_emb.T / self.temperature

        return {
            "drone_emb": drone_emb,
            "satellite_emb": satellite_emb,
            "logits": logits,
            "temperature": self.temperature,
            "drone_gate": self.drone_fusion.gate_value,
            "satellite_gate": self.satellite_fusion.gate_value,
        }


# ---------------------------------------------------------------------------
# Управление заморозкой параметров
# ---------------------------------------------------------------------------

def freeze_encoder(encoder: nn.Module) -> int:
    """Заморозить все параметры энкодера. Возвращает кол-во замороженных."""
    count = 0
    for p in encoder.parameters():
        p.requires_grad = False
        count += 1
    return count


def get_trainable_params(model: DualEncoder) -> dict[str, dict]:
    """Подсчёт обучаемых vs замороженных параметров по компонентам."""
    components = {
        "text_encoder": model.text_encoder,
        "drone_image_encoder": model.drone_image_encoder,
        "satellite_image_encoder": model.satellite_image_encoder,
        "drone_fusion": model.drone_fusion,
        "satellite_fusion": model.satellite_fusion,
        "drone_proj": model.drone_proj,
        "satellite_proj": model.satellite_proj,
    }
    result = {}
    for name, module in components.items():
        trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)
        total = sum(p.numel() for p in module.parameters())
        result[name] = {"trainable": trainable, "total": total}
    result["temperature"] = {"trainable": 1, "total": 1}
    return result


# ---------------------------------------------------------------------------
# Фабричная функция: собрать всё из чекпоинтов
# ---------------------------------------------------------------------------

def build_dual_encoder(
    dgtrs_checkpoint: str | Path,
    stripnet_checkpoint: str | Path,
    fused_dim: int = 512,
    shared_dim: int = 512,
    freeze_text: bool = True,
    freeze_image_backbone: bool = True,
    inject_mona: bool = True,
    mona_bottleneck: int = 64,
    init_gate: float = 0.7,
    baseline_mode: bool = False,
    device: str = "cpu",
) -> DualEncoder:
    """Собрать симметричный DualEncoder из чекпоинтов.

    Текстовый энкодер — ОДИН общий экземпляр, используется для текста
    и дрона, и спутника (это одна и та же языковая модель; разный текст
    подаётся на вход одного и того же набора весов).

    Визуальные энкодеры — ДВА отдельных экземпляра StripNetEncoder, по
    одному на домен (дрон / спутник), с отдельными Conv-MONA адаптерами.

    Слияние — ДВА отдельных GatedFusion (drone_fusion, satellite_fusion),
    с отдельными весами (проекции + скаляр α), так как оптимальное
    смешивание "картинка+текст" индивидуально для каждого визуального домена.
    Формула: σ(α)·img + (1 − σ(α))·text (первоисточник / протокол §3.2).

    Обучаются: оба GatedFusion, обе проекционные головы, оба набора
    Conv-MONA адаптеров, temperature.
    Замораживаются: текстовый энкодер целиком, backbone StripNet в обоих
    визуальных энкодерах.

    Args:
        dgtrs_checkpoint:       Путь к чекпоинту DGTRS-CLIP.
        stripnet_checkpoint:    Путь к pretrained StripNet (общий чекпоинт
                                 для инициализации обоих визуальных энкодеров).
        fused_dim:               Размерность вектора после слияния.
        shared_dim:              Размерность общего пространства сравнения.
        freeze_text:             Заморозить текстовый энкодер.
        freeze_image_backbone:   Заморозить backbone StripNet (в обоих).
        inject_mona:             Инжектировать Conv-MONA в оба энкодера.
        mona_bottleneck:         Bottleneck dim для Conv-MONA.
        init_gate:               Начальный вес картинки σ(α) в GatedFusion.
        baseline_mode:           Если True — gate ≡ 1.0 (B0 без текста).
        device:                  Устройство.
    """
    # 1. Текстовый энкодер (общий, заморожен, используется для обеих сторон)
    text_encoder = load_dgtrs_text_encoder(dgtrs_checkpoint, device="cpu")

    # 2. Два отдельных визуальных энкодера
    drone_image_encoder = StripNetEncoder(
        checkpoint_path=stripnet_checkpoint,
        out_dim=1024,
        load_pretrained=True,
    )
    satellite_image_encoder = StripNetEncoder(
        checkpoint_path=stripnet_checkpoint,
        out_dim=1024,
        load_pretrained=True,
    )

    # 3. Заморозка
    if freeze_text:
        n = freeze_encoder(text_encoder)
        LOGGER.info("❄️  Text encoder frozen: %d parameters (shared for drone+satellite)", n)

    if freeze_image_backbone:
        n1 = freeze_encoder(drone_image_encoder.backbone)
        n2 = freeze_encoder(satellite_image_encoder.backbone)
        LOGGER.info(
            "❄️  StripNet backbones frozen: drone=%d, satellite=%d parameters",
            n1, n2,
        )

    # 4. Conv-MONA адаптеры — отдельно для каждого визуального энкодера
    if inject_mona:
        from src.models.stripnet.conv_mona import inject_conv_mona_into_stripnet
        inject_conv_mona_into_stripnet(
            drone_image_encoder.backbone,
            bottleneck=mona_bottleneck,
            last_n_stages=2,
        )
        inject_conv_mona_into_stripnet(
            satellite_image_encoder.backbone,
            bottleneck=mona_bottleneck,
            last_n_stages=2,
        )

    # 5. Собрать DualEncoder (GatedFusion создаются внутри __init__)
    model = DualEncoder(
        text_encoder=text_encoder,
        drone_image_encoder=drone_image_encoder,
        satellite_image_encoder=satellite_image_encoder,
        fused_dim=fused_dim,
        shared_dim=shared_dim,
        init_gate=init_gate,
        baseline_mode=baseline_mode,
    )
    model = model.to(device)

    # 6. Логирование
    params = get_trainable_params(model)
    total_trainable = sum(v["trainable"] for v in params.values())
    total_all = sum(v["total"] for v in params.values())
    LOGGER.info(
        "🔗 DualEncoder (symmetric fusion) built: %s trainable / %s total params",
        f"{total_trainable:,}", f"{total_all:,}",
    )
    for name, info in params.items():
        LOGGER.info(
            "   %-26s %10s trainable / %10s total",
            name, f"{info['trainable']:,}", f"{info['total']:,}",
        )

    return model