from __future__ import annotations

import json
from functools import partial
from pathlib import Path
from typing import Callable


import torch
from timm.layers import to_2tuple
from timm.models.vision_transformer import VisionTransformer
from torch import Tensor, nn
from torchvision.transforms import ToTensor
from torchvision.transforms.v2 import Compose, Normalize, Resize


TARGET_CELL_SIZE = 40
LEMON_MOCO_ARCHITECTURES = ("vits8", "vitb8")


def build_lemon_moco_backbone(arch: str) -> nn.Module:
    if arch not in LEMON_MOCO_ARCHITECTURES:
        raise ValueError(
            f"Unknown LEMON MoCo architecture: {arch}. "
            f"Expected one of {LEMON_MOCO_ARCHITECTURES}."
        )
    patch_size = 8
    embed_dim = 384 if arch == "vits8" else 768
    model = VisionTransformer(
        img_size=TARGET_CELL_SIZE,
        patch_size=patch_size,
        embed_dim=embed_dim,
        depth=12,
        num_heads=12,
        mlp_ratio=4,
        qkv_bias=True,
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        num_classes=0,
    )
    return model



def load_lemon_moco_normalization() -> tuple[list[float], list[float]]:
    stats_path = Path(".") / "mean_std.json"
    if not stats_path.is_file():
        raise FileNotFoundError(f"LEMON MoCo normalization stats are missing: {stats_path}")
    with stats_path.open("r", encoding="utf-8") as stats_file:
        stats = json.load(stats_file)
    mean = [float(value) for value in stats["mean"]]
    std = [float(value) for value in stats["std"]]
    return mean, std


def build_lemon_moco_transform() -> Compose:
    mean, std = load_lemon_moco_normalization()
    transform = Compose(
        [
            ToTensor(),
            Resize(size=to_2tuple(TARGET_CELL_SIZE)),
            Normalize(mean=mean, std=std),
        ]
    )
    return transform


def extract_backbone_state_dict(checkpoint: dict) -> dict[str, Tensor]:
    if "state_dict" not in checkpoint:
        raise KeyError("LEMON MoCo checkpoint does not contain a 'state_dict' entry.")
    backbone_state_dict = {}
    for raw_key, value in checkpoint["state_dict"].items():
        key = raw_key.removeprefix("module.")
        if key.startswith("base_encoder.") and not key.startswith(
            ("base_encoder.head.", "base_encoder.fc.")
        ):
            backbone_state_dict[key.removeprefix("base_encoder.")] = value
    if not backbone_state_dict:
        raise ValueError("LEMON MoCo checkpoint did not contain base_encoder backbone weights.")
    return backbone_state_dict


def load_lemon_moco_model(arch: str = "vitb8") -> tuple[nn.Module, torch.dtype, Callable]:
    checkpoint_path =  Path(".") / f"lemon_{arch}.pth.tar"
    if not checkpoint_path.is_file():
        raise FileNotFoundError(f"LEMON MoCo checkpoint does not exist: {checkpoint_path}")
    transform = build_lemon_moco_transform()
    model = build_lemon_moco_backbone(arch)
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    backbone_state_dict = extract_backbone_state_dict(checkpoint)
    incompatible_keys = model.load_state_dict(backbone_state_dict, strict=False)
    if incompatible_keys.missing_keys:
        raise RuntimeError(
            f"Missing keys when loading LEMON MoCo backbone: {incompatible_keys.missing_keys}"
        )
    if incompatible_keys.unexpected_keys:
        raise RuntimeError(
            f"Unexpected keys when loading LEMON MoCo backbone: {incompatible_keys.unexpected_keys}"
        )
    return model, torch.float16, transform