"""Quickstart: predict, then correct aggregate shares for classifier error."""

from collections import Counter
import numpy as np, pandas as pd
from huggingface_hub import hf_hub_download

from predict import DemographicPredictor
from correction import correct_proportions

REPO = "Tijmen/age-gender-race-prediction"

# --- 1. per-image prediction -------------------------------------------------
p = DemographicPredictor()
print(p.predict("example.jpg"))

# --- 2. corpus shares --------------------------------------------------------
paths = ["a.jpg", "b.jpg", "c.jpg"]                     # your corpus
results = p.predict_batch(paths)

labels = ["White", "Black", "Asian", "Hispanic"]
counts = Counter(r["race_four"] for r in results)
n = sum(counts[l] for l in labels)
p_obs = np.array([counts[l] / n for l in labels]) if n else np.zeros(4)
print("observed :", dict(zip(labels, (p_obs * 100).round(2))))

# --- 3. correct for classifier error ----------------------------------------
# Photographic corpus -> the real-domain matrix. For AI-generated images use
# confusion_race_flux_ipw.csv instead; the error profiles differ.
M = pd.read_csv(hf_hub_download(REPO, "benchmark/confusion_race_real_ipw.csv"),
                index_col=0).values
print("corrected:", dict(zip(labels, (correct_proportions(p_obs, M) * 100).round(2))))