from typing import Optional

import os
import pandas as pd
import numpy as np
import tabpfn_client
from tabpfn_client import TabPFNClassifier
from tabpfn_client.constants import ModelVersion


def init_tabpfn(token: Optional[str] = None) -> None:
    """Configure tabpfn-client. Pass token explicitly or set TABPFN_API_TOKEN in the environment."""
    resolved = (token or "").strip() or os.environ.get("TABPFN_API_TOKEN", "").strip()
    if not resolved:
        raise RuntimeError("TabPFN API token is required.")
    tabpfn_client.set_access_token(resolved)

FEATURES = [
    "fixed_acidity", "volatile_acidity", "citric_acid", "residual_sugar",
    "chlorides", "free_sulfur_dioxide", "total_sulfur_dioxide", "density",
    "pH", "sulphates", "alcohol", "wine_type",
]

FEATURE_DISPLAY_NAMES = {
    "fixed_acidity": "Fixed Acidity",
    "volatile_acidity": "Volatile Acidity",
    "citric_acid": "Citric Acid",
    "residual_sugar": "Residual Sugar (g/L)",
    "chlorides": "Chlorides",
    "free_sulfur_dioxide": "Free SO₂ (mg/L)",
    "total_sulfur_dioxide": "Total SO₂ (mg/L)",
    "density": "Density (g/cm³)",
    "pH": "pH",
    "sulphates": "Sulphates",
    "alcohol": "Alcohol (%)",
    "wine_type": "Wine Type (0=Red, 1=White)",
}

QUALITY_LABELS = {0: "Low (≤5)", 1: "Medium (6)", 2: "High (≥7)"}
# Representative UCI quality scores (0–10) used to estimate an index from class probs.
QUALITY_INDEX_BY_BIN = {0: 4.5, 1: 6.0, 2: 7.5}


def quality_index_from_probs(probs: np.ndarray) -> float:
    return float(sum(probs[i] * QUALITY_INDEX_BY_BIN[i] for i in range(len(probs))))

RED_URL = "https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv"
WHITE_URL = "https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-white.csv"


def bin_quality(q: int) -> int:
    if q <= 5:
        return 0
    if q == 6:
        return 1
    return 2


def load_data() -> pd.DataFrame:
    red = pd.read_csv(RED_URL, sep=";")
    red["wine_type"] = 0
    white = pd.read_csv(WHITE_URL, sep=";")
    white["wine_type"] = 1
    df = pd.concat([red, white], ignore_index=True)
    df.columns = [c.strip().replace(" ", "_") for c in df.columns]
    df["quality_bin"] = df["quality"].apply(bin_quality)
    return df


def prepare_split(df: pd.DataFrame, test_size: int = 1000, random_state: int = 42):
    df = df.sample(frac=1, random_state=random_state).reset_index(drop=True)
    test_df = df.iloc[:test_size].reset_index(drop=True)
    train_df = df.iloc[test_size:].reset_index(drop=True)
    return train_df, test_df


def fit_model(train_df: pd.DataFrame) -> TabPFNClassifier:
    X = train_df[FEATURES].values.astype(float)
    y = train_df["quality_bin"].values
    clf = TabPFNClassifier.create_default_for_version(ModelVersion.V3)
    clf.fit(X, y)
    return clf


def predict(clf: TabPFNClassifier, features: dict) -> tuple[str, dict, float, int]:
    """Returns (predicted label, {label: probability}, estimated quality index, class index)."""
    row = np.array([[features[f] for f in FEATURES]], dtype=float)
    probs = clf.predict_proba(row)[0]
    pred_idx = int(np.argmax(probs))
    prob_map = {QUALITY_LABELS[i]: float(probs[i]) for i in range(len(probs))}
    quality_index = round(quality_index_from_probs(probs), 1)
    return QUALITY_LABELS[pred_idx], prob_map, quality_index, pred_idx