from __future__ import annotations

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


class Bridge(nn.Module):
    def __init__(
        self,
        vision_dim: int,
        text_dim: int,
        hidden_dim: int = 1536,
        linear: bool = False,
    ):
        super().__init__()
        self.vision_dim = vision_dim
        self.text_dim = text_dim
        self.hidden_dim = hidden_dim
        self.linear = linear
        if linear:
            self.network = nn.Linear(vision_dim, text_dim, bias=False)
        else:
            self.network = nn.Sequential(
                nn.LayerNorm(vision_dim),
                nn.Linear(vision_dim, hidden_dim),
                nn.GELU(),
                nn.Linear(hidden_dim, text_dim),
            )

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

    def config(self) -> dict:
        return {
            "vision_dim": self.vision_dim,
            "text_dim": self.text_dim,
            "hidden_dim": self.hidden_dim,
            "linear": self.linear,
        }


class PrefixAdapter(nn.Module):
    def __init__(
        self,
        semantic_dim: int,
        lm_dim: int,
        prefix_length: int = 8,
        hidden_dim: int = 2048,
    ):
        super().__init__()
        self.semantic_dim = semantic_dim
        self.lm_dim = lm_dim
        self.prefix_length = prefix_length
        self.hidden_dim = hidden_dim
        self.network = nn.Sequential(
            nn.LayerNorm(semantic_dim),
            nn.Linear(semantic_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, prefix_length * lm_dim),
        )

    def forward(self, semantic: torch.Tensor) -> torch.Tensor:
        prefix = self.network(semantic.float())
        return prefix.view(-1, self.prefix_length, self.lm_dim)

    def config(self) -> dict:
        return {
            "semantic_dim": self.semantic_dim,
            "lm_dim": self.lm_dim,
            "prefix_length": self.prefix_length,
            "hidden_dim": self.hidden_dim,
        }


def load_bridge(path: str, device: str = "cpu") -> tuple[Bridge, dict]:
    state = torch.load(path, map_location="cpu", weights_only=False)
    model = Bridge(**state["config"])
    model.load_state_dict(state["state_dict"])
    model.to(device).eval()
    return model, state


def load_prefix(path: str, device: str = "cpu") -> tuple[PrefixAdapter, dict]:
    state = torch.load(path, map_location="cpu", weights_only=False)
    model = PrefixAdapter(**state["config"])
    model.load_state_dict(state["state_dict"])
    model.to(device).eval()
    return model, state