"""
FBP5500 percentile rater — minimal version.

Flow: pick gender -> upload -> 15 pairwise comparisons -> bell curve + percentile.
"""

import numpy as np
import streamlit as st
import matplotlib.pyplot as plt
from PIL import Image
from datasets import load_dataset, concatenate_datasets
from scipy.stats import gaussian_kde

HF_DATASET = "MnLgt/scut-fbp5500"
N_BINS = 100
SCORE_MIN, SCORE_MAX = 1.0, 5.0
TAU = 0.45
EPSILON_UNIFORM = 0.08
SHARPNESS = 4              # exponent on info-gain weights; higher = more concentrated near posterior mean
TARGET_COMPARISONS = 15
DISPLAY_SIZE = (350, 350)  # FBP5500 images are 350x350; uploaded photo resized to match

# ---- data ------------------------------------------------------------------

@st.cache_resource(show_spinner="Loading dataset...")
def load_ds():
    ds = load_dataset(HF_DATASET)
    return concatenate_datasets([ds["train"], ds["test"]]) if "test" in ds else ds["train"]

@st.cache_data(show_spinner=False)
def get_meta(_fp: str):
    """Return (scores, genders, races) as numpy arrays aligned with dataset indices."""
    ds = load_ds()
    df = ds.remove_columns(["image"]).to_pandas()
    return np.array(df["beauty_score"]), np.array(df["gender"]), np.array(df["race"])

# ---- bayesian --------------------------------------------------------------

def bin_centers():
    edges = np.linspace(SCORE_MIN, SCORE_MAX, N_BINS + 1)
    return 0.5 * (edges[:-1] + edges[1:])

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))

def update(post, ref_score, said_more):
    p = sigmoid((bin_centers() - ref_score) / TAU)
    new = post * (p if said_more else (1 - p))
    s = new.sum()
    return new / s if s > 0 else post

def info_weights(post, scores):
    diffs = (bin_centers()[None, :] - scores[:, None]) / TAU
    p = sigmoid(diffs)
    e = (p * post[None, :]).sum(axis=1)
    return e * (1 - e)

def sample_idx(pool_indices, all_scores, post, seen, rng):
    """Sample from pool_indices (dataset indices), excluding seen."""
    avail = np.array([i for i in pool_indices if i not in seen])
    if len(avail) == 0:
        avail = pool_indices
    w = info_weights(post, all_scores[avail])
    w = w ** SHARPNESS
    w = (1 - EPSILON_UNIFORM) * (w / w.sum()) + EPSILON_UNIFORM / len(w)
    return int(rng.choice(avail, p=w))

# ---- image helpers ---------------------------------------------------------

def fit_to_size(img: Image.Image, size=DISPLAY_SIZE) -> Image.Image:
    """Resize image to fit within `size` keeping aspect ratio, then pad to exact size."""
    img = img.copy()
    img.thumbnail(size, Image.LANCZOS)
    # pad with neutral background to exact size
    canvas = Image.new("RGB", size, (240, 240, 240))
    x = (size[0] - img.width) // 2
    y = (size[1] - img.height) // 2
    canvas.paste(img, (x, y))
    return canvas

# ---- UI --------------------------------------------------------------------

st.set_page_config(page_title="Face Rater", layout="wide")

st.markdown("""
<style>
    #MainMenu, footer, header {visibility: hidden;}
    .block-container {padding-top: 2rem; max-width: 900px;}
    .stButton > button {width: 100%; padding: 0.75rem; font-size: 1rem; font-weight: 500;}
    div[data-testid="stImage"] img {border-radius: 8px;}
</style>
""", unsafe_allow_html=True)

ds = load_ds()
all_scores, all_genders, all_races = get_meta(ds._fingerprint)

# session init
if "posterior" not in st.session_state:
    st.session_state.posterior = np.ones(N_BINS) / N_BINS
    st.session_state.seen = set()
    st.session_state.n_done = 0
    st.session_state.current_idx = None
    st.session_state.rng = np.random.default_rng()
    st.session_state.done = False
    st.session_state.gender = None
    st.session_state.race = None
    st.session_state.pool_indices = None

# step 1: gender pick
if st.session_state.gender is None:
    st.title("Face percentile rater")
    st.write("Compare against which group?")
    c1, c2 = st.columns(2)
    if c1.button("Men", use_container_width=True):
        st.session_state.gender = "Male"
        st.rerun()
    if c2.button("Women", use_container_width=True):
        st.session_state.gender = "Female"
        st.rerun()
    st.stop()

# step 1b: race pick
if st.session_state.race is None:
    st.title("Face percentile rater")
    st.write(f"Selected: {st.session_state.gender.lower()}s")
    st.write("Compare against which race?")
    c1, c2, c3 = st.columns(3)
    if c1.button("Caucasian", use_container_width=True):
        st.session_state.race = "Caucasian"
        st.rerun()
    if c2.button("Asian", use_container_width=True):
        st.session_state.race = "Asian"
        st.rerun()
    if c3.button("Both", use_container_width=True):
        st.session_state.race = "Both"
        st.rerun()
    st.stop()

# precompute pool indices for the chosen gender + race
if st.session_state.pool_indices is None:
    gender_mask = all_genders == st.session_state.gender
    if st.session_state.race == "Both":
        mask = gender_mask
    else:
        mask = gender_mask & (all_races == st.session_state.race)
    st.session_state.pool_indices = np.where(mask)[0]

# step 2: upload OR pick a calibration face
if "uploaded_img" not in st.session_state:
    st.title("Face percentile rater")
    race_label = "Asian + Caucasian" if st.session_state.race == "Both" else st.session_state.race
    st.caption(f"Comparing against: {race_label} {st.session_state.gender.lower()}s "
               f"({len(st.session_state.pool_indices)} reference faces)")

    uploaded = st.file_uploader("Upload a face photo", type=["jpg", "jpeg", "png", "webp"])
    if uploaded is not None:
        st.session_state.uploaded_img = Image.open(uploaded).convert("RGB")
        # mark that this came from upload (so we don't exclude it from sampling)
        st.session_state.calibration_idx = None
        st.rerun()

    st.markdown("**Or test the rater with a known face from the dataset:**")
    st.caption("Picks a face from FBP5500 near the chosen percentile. See if the rater recovers it.")

    pool_scores = all_scores[st.session_state.pool_indices]
    sorted_pool_scores = np.sort(pool_scores)
    percentiles = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]

    # 11 buttons in a row
    cols = st.columns(11)
    for col, p in zip(cols, percentiles):
        if col.button(str(p), key=f"pct_{p}", use_container_width=True):
            # find the score at the chosen percentile within the gender pool
            target_score = float(np.percentile(pool_scores, p))
            # pick a face whose score is closest to target (with a small random jitter
            # among the k nearest so repeated clicks don't always show the same face)
            pool_idx_array = st.session_state.pool_indices
            scores_in_pool = all_scores[pool_idx_array]
            distances = np.abs(scores_in_pool - target_score)
            k_nearest = np.argsort(distances)[:5]  # 5 closest faces
            chosen_local = st.session_state.rng.choice(k_nearest)
            chosen_idx = int(pool_idx_array[chosen_local])

            st.session_state.uploaded_img = ds[chosen_idx]["image"].convert("RGB")
            st.session_state.calibration_idx = chosen_idx
            st.session_state.calibration_true_score = float(ds[chosen_idx]["beauty_score"])
            # exclude this face from being shown as a reference
            st.session_state.seen.add(chosen_idx)
            st.rerun()

    st.stop()

uploaded_img = st.session_state.uploaded_img
uploaded_display = fit_to_size(uploaded_img)

# step 4: end screen (KDE on the filtered pool)
if st.session_state.done:
    centers = bin_centers()
    theta_hat = float((st.session_state.posterior * centers).sum())
    pool_scores = all_scores[st.session_state.pool_indices]
    sorted_scores = np.sort(pool_scores)
    pct = float(np.searchsorted(sorted_scores, theta_hat, side="right") / len(sorted_scores) * 100)

    kde = gaussian_kde(pool_scores, bw_method=0.3)
    xs = np.linspace(1, 5, 400)
    ys = kde(xs)

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.fill_between(xs, ys, color="#cdd9e5", alpha=0.6)
    ax.plot(xs, ys, color="#5a7184", linewidth=2)

    y_at = float(kde(theta_hat)[0])
    ax.plot([theta_hat, theta_hat], [0, y_at], color="#c0392b", linewidth=2.5)
    ax.scatter([theta_hat], [y_at], color="#c0392b", s=120, zorder=5)
    ax.annotate(f"  you ({pct:.0f}th pct)",
                xy=(theta_hat, y_at),
                xytext=(theta_hat + 0.05, y_at + max(ys) * 0.05),
                fontsize=13, color="#c0392b", fontweight="bold")

    # if this was a calibration face, plot the ground truth too
    cal_score = st.session_state.get("calibration_true_score")
    if cal_score is not None:
        true_pct = float(np.searchsorted(sorted_scores, cal_score, side="right") / len(sorted_scores) * 100)
        y_true = float(kde(cal_score)[0])
        ax.plot([cal_score, cal_score], [0, y_true], color="#27ae60", linewidth=2.5, linestyle="--")
        ax.scatter([cal_score], [y_true], color="#27ae60", s=120, zorder=5, marker="D")
        ax.annotate(f"  truth ({true_pct:.0f}th pct)",
                    xy=(cal_score, y_true),
                    xytext=(cal_score + 0.05, y_true + max(ys) * 0.12),
                    fontsize=13, color="#27ae60", fontweight="bold")

    ax.set_xlim(1, 5)
    ax.set_ylim(0, max(ys) * 1.2)
    race_label = "Asian + Caucasian" if st.session_state.race == "Both" else st.session_state.race
    ax.set_xlabel(f"FBP5500 beauty score ({race_label} {st.session_state.gender.lower()}s)", fontsize=12)
    ax.set_yticks([])
    for spine in ("top", "right", "left"):
        ax.spines[spine].set_visible(False)
    ax.set_title(f"{pct:.0f}th percentile  (θ ≈ {theta_hat:.2f})", fontsize=16, pad=15)

    st.pyplot(fig)

    c1, c2 = st.columns(2)
    with c1:
        st.image(uploaded_display, caption="Your photo", use_container_width=True)
    with c2:
        if st.button("Start over", use_container_width=True):
            for k in list(st.session_state.keys()):
                del st.session_state[k]
            st.rerun()
    st.stop()

# step 3: comparison loop
if st.session_state.current_idx is None:
    st.session_state.current_idx = sample_idx(
        st.session_state.pool_indices, all_scores,
        st.session_state.posterior, st.session_state.seen, st.session_state.rng
    )

cur = st.session_state.current_idx
row = ds[cur]
ref_score = float(row["beauty_score"])
ref_display = fit_to_size(row["image"])

col_l, col_r = st.columns(2, gap="medium")
with col_l:
    st.image(uploaded_display, use_container_width=True)
    pick_left = st.button("This one", key="pick_left", use_container_width=True)
with col_r:
    st.image(ref_display, use_container_width=True)
    pick_right = st.button("This one", key="pick_right", use_container_width=True)

def advance(said_more: bool):
    st.session_state.posterior = update(st.session_state.posterior, ref_score, said_more)
    st.session_state.seen.add(cur)
    st.session_state.n_done += 1
    st.session_state.current_idx = None
    if st.session_state.n_done >= TARGET_COMPARISONS:
        st.session_state.done = True

if pick_left:
    advance(True); st.rerun()
if pick_right:
    advance(False); st.rerun()

remaining = TARGET_COMPARISONS - st.session_state.n_done
st.divider()
st.progress(st.session_state.n_done / TARGET_COMPARISONS,
            text=f"{remaining} more rating{'s' if remaining != 1 else ''} to go")