"""
Simple emotion classification with a Deit-Base model hosted on the Hugging Face Hub.

Install:
    pip install torch timm safetensors huggingface_hub pillow

Quick use:
    from emotion_classifier import EmotionClassifier

    clf = EmotionClassifier()
    result = clf.predict("path/to/image.jpg")
    print(result)
    # {'prediction': 'Neutral', 'confidence': 0.95, 'Optimistic': 0.01,
    #  'Pessimistic': 0.01, 'Hostile': 0.02, 'Neutral': 0.95}

Command line use:
    python emotion_classifier.py path/to/image_or_folder
    python emotion_classifier.py path/to/folder --output results.csv
"""

import argparse
from pathlib import Path

import torch
import torch.nn.functional as F
import timm
from PIL import Image
from torchvision import transforms
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

# ---------------------------------------------------------------------------
# Fixed model settings (must match how the checkpoint was trained)
# ---------------------------------------------------------------------------
REPO_ID       = "Simon14142/emotion-recognition-4class-deit"
FILENAME      = "DeiT-Base_checkpoint_final.safetensors"
MODEL_NAME    = "deit_base_patch16_224"
IMG_SIZE      = 224
EMOTION_NAMES = ["Optimistic", "Pessimistic", "Hostile", "Neutral"]
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff"}


class EmotionClassifier:
    """Loads the Deit-Base emotion model straight from the Hugging Face Hub."""

    def __init__(self, repo_id: str = REPO_ID, filename: str = FILENAME,
                 device: str | None = None, precision: str = "fp32"):
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))

        checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)

        self.model = timm.create_model(MODEL_NAME, pretrained=False, num_classes=len(EMOTION_NAMES))
        state_dict = load_file(checkpoint_path, device="cpu")
        missing, unexpected = self.model.load_state_dict(state_dict, strict=True)
        if missing:
            print(f"Warning - missing keys: {missing}")
        if unexpected:
            print(f"Warning - unexpected keys: {unexpected}")

        self.precision = precision
        if precision == "bf16":
            if self.device.type == "cuda" and torch.cuda.is_bf16_supported():
                self.model = self.model.to(torch.bfloat16)
            else:
                print("BF16 not supported here, falling back to FP32.")
                self.precision = "fp32"

        self.model = self.model.to(self.device).eval()

        self.transform = transforms.Compose([
            transforms.Resize((IMG_SIZE, IMG_SIZE)),
            transforms.ToTensor(),
            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
        ])

    @torch.no_grad()
    def predict(self, image, threshold: float = 0.4) -> dict:
        """image: a file path (str/Path) or a PIL.Image. Returns a result dict."""
        if not isinstance(image, Image.Image):
            image = Image.open(image).convert("RGB")

        tensor = self.transform(image).unsqueeze(0).to(self.device)
        if self.precision == "bf16":
            tensor = tensor.to(torch.bfloat16)

        logits = self.model(tensor)
        probs = F.softmax(logits.float(), dim=1).squeeze(0).cpu()
        pred_idx = int(probs.argmax())
        conf = float(probs[pred_idx])
        label = EMOTION_NAMES[pred_idx] if conf >= threshold else "Unknown"

        return {
            "prediction": label,
            "confidence": round(conf, 4),
            **{name: round(float(probs[i]), 4) for i, name in enumerate(EMOTION_NAMES)},
        }


def _collect_image_paths(input_path: str) -> list[str]:
    p = Path(input_path)
    if p.is_file():
        return [str(p)]
    if p.is_dir():
        paths = sorted(str(f) for f in p.iterdir() if f.suffix.lower() in IMAGE_EXTENSIONS)
        if not paths:
            raise FileNotFoundError(f"No images found in: {input_path}")
        return paths
    raise FileNotFoundError(f"Path does not exist: {input_path}")


def main():
    parser = argparse.ArgumentParser(description="Emotion classification with Deit-Base")
    parser.add_argument("input", help="Path to an image file or a folder of images")
    parser.add_argument("--threshold", type=float, default=0.4)
    parser.add_argument("--output", help="Optional CSV path to save results")
    args = parser.parse_args()

    clf = EmotionClassifier()
    paths = _collect_image_paths(args.input)

    rows = []
    for path in paths:
        result = clf.predict(path, threshold=args.threshold)
        result["image_path"] = path
        rows.append(result)
        print(f"{Path(path).name:<40} -> {result['prediction']:<12} ({result['confidence']:.3f})")

    if args.output:
        import csv
        with open(args.output, "w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=["image_path", "prediction", "confidence", *EMOTION_NAMES])
            writer.writeheader()
            writer.writerows(rows)
        print(f"\nSaved {len(rows)} results to {args.output}")


if __name__ == "__main__":
    main()