import json

import cv2
import numpy as np
import timm
import torch
import torch.nn as nn
from PIL import Image

# Fixed parameters, no state mutated during apply() - safe to share
_CLAHE = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))

# Deterministic views - all valid for fundus images (no natural orientation)
_TTA_AUGMENTS = [
    lambda t: t,
    lambda t: torch.flip(t, dims=[-1]),
    lambda t: torch.flip(t, dims=[-2]),
    lambda t: torch.rot90(t, k=1, dims=[-2, -1]),
]


class DualEyeModel(nn.Module):
    def __init__(self, backbone: str, num_classes: int, dropout: float):
        super().__init__()
        # pretrained=False: the fine-tuned state dict is loaded on top, so
        # downloading ImageNet weights here would be wasted bandwidth
        self.backbone = timm.create_model(backbone, pretrained=False, num_classes=0)
        features = self.backbone.num_features
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(features * 2, num_classes),
        )

    def forward(self, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
        both = torch.cat([left, right], dim=0)
        features = self.backbone(both)
        left_feat, right_feat = features.chunk(2, dim=0)
        return self.classifier(torch.cat([left_feat, right_feat], dim=1))


class RetinalEnsemble(nn.Module):
    def __init__(self, config: dict):
        super().__init__()
        self.labels = list(config["labels"])
        self.image_size = int(config["image_size"])
        self.models = nn.ModuleList(
            DualEyeModel(b, config["num_classes"], config["dropout"])
            for b in config["backbones"]
        )
        # Each backbone expects its own input normalization; applied here so a
        # single un-normalized [0, 1] tensor feeds the whole ensemble
        for i, (mean, std) in enumerate(zip(config["norm_means"], config["norm_stds"])):
            self.register_buffer(f"mean_{i}", torch.tensor(mean, dtype=torch.float32).view(3, 1, 1))
            self.register_buffer(f"std_{i}", torch.tensor(std, dtype=torch.float32).view(3, 1, 1))
        self.register_buffer(
            "thresholds", torch.tensor(config["thresholds"], dtype=torch.float32)
        )

    def forward(self, left: torch.Tensor, right: torch.Tensor, tta: bool = False) -> torch.Tensor:
        augments = _TTA_AUGMENTS if tta else _TTA_AUGMENTS[:1]
        views = []
        for aug in augments:
            l, r = aug(left), aug(right)
            probs = []
            for i, model in enumerate(self.models):
                mean = getattr(self, f"mean_{i}")
                std = getattr(self, f"std_{i}")
                probs.append(torch.sigmoid(model((l - mean) / std, (r - mean) / std)))
            views.append(torch.stack(probs).mean(dim=0))
        return torch.stack(views).mean(dim=0)

    def preprocess(self, image) -> torch.Tensor:
        if not isinstance(image, Image.Image):
            image = Image.open(image)
        image = image.convert("RGB")
        lab = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2LAB)
        lab[:, :, 0] = _CLAHE.apply(lab[:, :, 0])
        image = Image.fromarray(cv2.cvtColor(lab, cv2.COLOR_LAB2RGB))
        image = image.resize((self.image_size, self.image_size), Image.BILINEAR)
        arr = np.asarray(image, dtype=np.float32) / 255.0
        return torch.from_numpy(arr).permute(2, 0, 1).contiguous()

    @torch.no_grad()
    def predict(self, left_image, right_image, tta: bool = True) -> tuple[dict, dict]:
        device = self.thresholds.device
        left = self.preprocess(left_image).unsqueeze(0).to(device)
        right = self.preprocess(right_image).unsqueeze(0).to(device)
        probs = self.forward(left, right, tta=tta)[0]
        preds = probs >= self.thresholds
        probabilities = {label: round(float(probs[i]), 4) for i, label in enumerate(self.labels)}
        predictions = {label: bool(preds[i]) for i, label in enumerate(self.labels)}
        return probabilities, predictions


def load_ensemble_from_hub(repo_id: str, device: str = "cpu", revision: str | None = None) -> RetinalEnsemble:
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    config_path = hf_hub_download(repo_id, "config.json", revision=revision)
    with open(config_path) as f:
        config = json.load(f)
    weights_path = hf_hub_download(repo_id, config["weights_file"], revision=revision)

    model = RetinalEnsemble(config).to(device)
    model.load_state_dict(load_file(weights_path, device=device))
    model.eval()
    return model