import json
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from torch import nn


class Projector(nn.Module):
    def __init__(self, config: dict):
        super().__init__()
        self.spatial_merge_size = 2
        self.linear_1 = nn.Linear(config["input_size"], config["hidden_size"], bias=True)
        self.linear_2 = nn.Linear(config["hidden_size"], config["output_size"], bias=True)
        self.merge_linear_1 = nn.Linear(config["merged_hidden_size"], config["hidden_size"], bias=True)
        self.merge_linear_2 = nn.Linear(config["hidden_size"], config["output_size"], bias=True)

    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
        hidden = torch.nn.functional.gelu(self.linear_1(image_features))
        hidden = self.linear_2(hidden)
        hidden = hidden.reshape(hidden.shape[0] // (self.spatial_merge_size ** 2), -1)
        return self.merge_linear_2(torch.nn.functional.gelu(self.merge_linear_1(hidden)))


def load_projector(model: str | Path) -> Projector:
    path = Path(model)
    if path.is_dir():
        config_path = path / "projector_config.json"
        weights_path = path / "projector.safetensors"
    else:
        config_path = Path(hf_hub_download(str(model), "projector_config.json"))
        weights_path = Path(hf_hub_download(str(model), "projector.safetensors"))
    projector = Projector(json.loads(config_path.read_text()))
    projector.load_state_dict(load_file(weights_path))
    return projector