from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from transformers import SiglipVisionConfig, SiglipVisionModel

ROOT = Path(__file__).resolve().parents[1]
MODELS_DIR = ROOT / "models"


class NutritionHead(nn.Module):
    def __init__(self, input_dim: int = 768) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(input_dim),
            nn.Linear(input_dim, 768),
            nn.GELU(),
            nn.Dropout(0.15),
            nn.Linear(768, 256),
            nn.GELU(),
            nn.Dropout(0.10),
            nn.Linear(256, 5),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.softplus(self.net(x))


@dataclass(frozen=True)
class CalorieKPrediction:
    kcal: float
    mass_g: float
    protein_g: float
    carbohydrate_g: float
    fat_g: float
    interval_90_radius_kcal: float
    model: str
    device: str


def prediction_to_dict(prediction: CalorieKPrediction) -> dict[str, Any]:
    radius = prediction.interval_90_radius_kcal
    return {
        "kcal": round(prediction.kcal),
        "mass_g": round(prediction.mass_g),
        "protein_g": round(prediction.protein_g, 1),
        "carbs_g": round(prediction.carbohydrate_g, 1),
        "fat_g": round(prediction.fat_g, 1),
        "interval_90_kcal": round(radius),
        "kcal_low": max(0, round(prediction.kcal - radius)),
        "kcal_high": round(prediction.kcal + radius),
        "model": prediction.model,
        "device": prediction.device,
        "notes": "",
    }


def _load_json(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def preprocess_image(image: Image.Image, image_size: int = 224) -> torch.Tensor:
    rgb = image.convert("RGB").resize((image_size, image_size), Image.Resampling.BICUBIC)
    array = np.asarray(rgb, dtype=np.float32) / 255.0
    pixels = torch.from_numpy(array).permute(2, 0, 1)
    return ((pixels - 0.5) / 0.5).unsqueeze(0)


def ensure_weights() -> Path:
    checkpoint_path = MODELS_DIR / "config.json"
    config = _load_json(checkpoint_path)
    weights = MODELS_DIR / config["checkpoint_file"]
    if weights.exists():
        return weights
    from huggingface_hub import hf_hub_download

    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    hf_hub_download(
        repo_id="Dralkh/CalorieK",
        filename=config["checkpoint_file"],
        local_dir=str(MODELS_DIR),
    )
    return weights


class CalorieK:
    def __init__(self) -> None:
        self.config = _load_json(MODELS_DIR / "config.json")
        calibration = _load_json(MODELS_DIR / self.config["calibration_file"])
        self.interval = float(
            self.config.get("interval_90_radius_kcal") or calibration.get("radius_kcal") or 118.95
        )
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.image_size = int(self.config.get("image_size", 224))

        checkpoint_path = ensure_weights()

        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
        vision_cfg = SiglipVisionConfig.from_json_file(str(MODELS_DIR / "vision_config.json"))
        self.vision = SiglipVisionModel(vision_cfg)
        vision_target = self.vision.vision_model if hasattr(self.vision, "vision_model") else self.vision
        vision_target.load_state_dict(checkpoint["vision_state"])

        self.head = NutritionHead(int(self.config.get("hidden_size", 768)))
        self.head.load_state_dict(checkpoint["head_state"])
        self.vision.to(self.device).eval()
        self.head.to(self.device).eval()

    @torch.inference_mode()
    def predict(self, image: Image.Image) -> CalorieKPrediction:
        pixels = preprocess_image(image, self.image_size).to(self.device)
        out = self.vision(pixel_values=pixels)
        feats = out.pooler_output if out.pooler_output is not None else out.last_hidden_state[:, 0]
        kcal, mass_g, protein_g, carb_g, fat_g = self.head(feats.float())[0].tolist()
        return CalorieKPrediction(
            kcal=float(kcal),
            mass_g=float(mass_g),
            protein_g=float(protein_g),
            carbohydrate_g=float(carb_g),
            fat_g=float(fat_g),
            interval_90_radius_kcal=self.interval,
            model="plato",
            device=str(self.device),
        )


_model: CalorieK | None = None


def get_model() -> CalorieK:
    global _model
    if _model is None:
        _model = CalorieK()
    return _model