from __future__ import annotations

"""Dual encoder for caption quality test on cross-view geo-localization.

GeoRSCLIP ViT-B/32 backbone (image + text towers, shared 512-dim space).
Image encoder frozen. Text encoder with partial unfreeze.

Architecture:
  Query  branch: GeoRSCLIP_img(drone) + GeoRSCLIP_text(caption) -> GatedFusion -> proj -> query_emb
  Gallery branch: GeoRSCLIP_img(sat) -> proj -> gallery_emb
  Loss: InfoNCE(query_emb, gallery_emb)

Baseline mode: fusion gate forced to 1.0 (text ignored).
"""

from typing import Literal

import gin
import open_clip
import torch
import torch.nn as nn
import torch.nn.functional as F


class ProjectionHead(nn.Module):
    """MLP projection head with L2 normalization."""

    def __init__(
        self,
        in_dim: int = 512,
        out_dim: int = 512,
        use_mlp: bool = False,
        hidden_dim: int | None = None,
    ) -> None:
        super().__init__()
        if use_mlp:
            hidden_dim = hidden_dim or (2 * in_dim)
            self.proj = nn.Sequential(
                nn.Linear(in_dim, hidden_dim),
                nn.GELU(),
                nn.Linear(hidden_dim, out_dim),
            )
        else:
            self.proj = nn.Linear(in_dim, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.normalize(self.proj(x), dim=-1)


@gin.configurable
class GatedFusion(nn.Module):
    """Learnable gated fusion of image and text embeddings.

    q = sigma(alpha) * img + (1 - sigma(alpha)) * text

    alpha is a single learnable scalar, initialized so that gate ~ init_gate.
    When baseline_mode=True, gate is clamped to 1.0 (text contribution = 0).
    """

    def __init__(self, init_gate: float = 0.7, baseline_mode: bool = False) -> None:
        super().__init__()
        # alpha is in logit space: sigmoid(alpha) = 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

    def forward(
        self,
        img_feat: torch.Tensor,
        text_feat: torch.Tensor | None,
    ) -> torch.Tensor:
        if text_feat is None or self.baseline_mode:
            return img_feat
        gate = torch.sigmoid(self.alpha)
        return gate * img_feat + (1.0 - gate) * text_feat

    @property
    def gate_value(self) -> float:
        """Current gate value (image weight). 1.0 = text ignored."""
        if self.baseline_mode:
            return 1.0
        return torch.sigmoid(self.alpha).item()


@gin.configurable
class DualEncoderCaptionTest(nn.Module):
    """GeoRSCLIP dual encoder with gated text fusion on query branch.

    Args:
        variant: open_clip model variant name.
        pretrained_path: Path to GeoRSCLIP checkpoint.
        unfreeze_mode: Text encoder unfreeze strategy.
        embed_dim: Output retrieval embedding dimension.
        use_mlp_heads: Use 2-layer MLP projection heads.
        baseline_mode: If True, fusion gate = 1.0 (no text).
        init_gate: Initial gate value (image weight).
        device: torch device.
    """

    def __init__(
        self,
        variant: str = "ViT-B-32",
        pretrained_path: str = "RS5M_ViT-B-32.pt",
        unfreeze_mode: Literal["none", "projection", "last_block"] = "last_block",
        embed_dim: int = 512,
        use_mlp_heads: bool = False,
        baseline_mode: bool = False,
        init_gate: float = 0.7,
        device: str = "cuda",
    ) -> None:
        super().__init__()
        self.embed_dim = embed_dim
        self.device = device
        self.baseline_mode = baseline_mode

        # Load GeoRSCLIP via open_clip.
        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
            model_name=variant,
            pretrained=pretrained_path,
            device=device,
        )
        self.tokenizer = open_clip.get_tokenizer(variant)

        native_dim = self._infer_native_dim()

        # Freeze everything.
        for p in self.model.parameters():
            p.requires_grad = False

        # Selectively unfreeze text encoder (only if not baseline).
        if not baseline_mode:
            self._apply_unfreeze(unfreeze_mode)

        # Gated fusion on query branch.
        self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)

        # Projection heads.
        self.proj_query = ProjectionHead(
            in_dim=native_dim, out_dim=embed_dim, use_mlp=use_mlp_heads,
        )
        self.proj_gallery = ProjectionHead(
            in_dim=native_dim, out_dim=embed_dim, use_mlp=use_mlp_heads,
        )

    def _infer_native_dim(self) -> int:
        if hasattr(self.model, "text_projection"):
            shape = self.model.text_projection.shape
            return int(shape[1] if shape.ndim == 2 else shape[0])
        return 512

    def _apply_unfreeze(
        self,
        unfreeze_mode: Literal["none", "projection", "last_block"],
    ) -> None:
        if unfreeze_mode == "none":
            return
        # Unfreeze text_projection.
        if hasattr(self.model, "text_projection"):
            tp = self.model.text_projection
            if isinstance(tp, nn.Parameter):
                tp.requires_grad = True
            elif isinstance(tp, nn.Module):
                for p in tp.parameters():
                    p.requires_grad = True
        # Unfreeze last transformer block.
        if unfreeze_mode == "last_block" and hasattr(self.model, "transformer"):
            for p in self.model.transformer.resblocks[-1].parameters():
                p.requires_grad = True

    def encode_image(self, images: torch.Tensor) -> torch.Tensor:
        feats = self.model.encode_image(images)
        return F.normalize(feats, dim=-1)

    def encode_text(self, texts: list[str]) -> torch.Tensor:
        tokens = self.tokenizer(list(texts)).to(self.device).long()
        feats = self.model.encode_text(tokens)
        return F.normalize(feats, dim=-1)

    def forward(
        self,
        drone_img: torch.Tensor,
        sat_img: torch.Tensor,
        caption_drone: list[str] | None = None,
    ) -> dict[str, torch.Tensor]:
        """Forward pass.

        Args:
            drone_img: Drone images [B, 3, H, W].
            sat_img: Satellite images [B, 3, H, W].
            caption_drone: Drone captions (P3 fingerprint), one per sample.

        Returns:
            Dict with 'query' [B, embed_dim], 'gallery' [B, embed_dim],
            and 'gate' (scalar) for logging.
        """
        # Gallery branch: satellite only.
        sat_feat = self.encode_image(sat_img)
        gallery = self.proj_gallery(sat_feat)

        # Query branch: drone + optional text fusion.
        drone_feat = self.encode_image(drone_img)

        text_feat = None
        if caption_drone is not None and not self.baseline_mode:
            text_feat = self.encode_text(caption_drone)

        fused = self.fusion(drone_feat, text_feat)
        query = self.proj_query(fused)

        return {
            "query": query,
            "gallery": gallery,
            "gate": self.fusion.gate_value,
        }

    def trainable_parameters(self) -> list[nn.Parameter]:
        return [p for p in self.parameters() if p.requires_grad]
