"""Inference helper for the UCI Covertype decision-tree classifier."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Mapping, Sequence

import joblib
import pandas as pd


ROOT = Path(__file__).resolve().parent
FEATURE_ORDER = [
    "Elevation", "Aspect", "Slope", "Horizontal_Distance_To_Hydrology",
    "Vertical_Distance_To_Hydrology", "Horizontal_Distance_To_Roadways",
    "Hillshade_9am", "Hillshade_Noon", "Hillshade_3pm",
    "Horizontal_Distance_To_Fire_Points",
    *[f"Wilderness_Area{i}" for i in range(1, 5)],
    *[f"Soil_Type{i}" for i in range(1, 41)],
]
LABELS = {
    1: "Spruce/Fir",
    2: "Lodgepole Pine",
    3: "Ponderosa Pine",
    4: "Cottonwood/Willow",
    5: "Aspen",
    6: "Douglas-fir",
    7: "Krummholz",
}


def _as_frame(record: Mapping[str, Any] | Sequence[Any] | pd.DataFrame) -> pd.DataFrame:
    if isinstance(record, pd.DataFrame):
        missing = [name for name in FEATURE_ORDER if name not in record.columns]
        if missing:
            raise ValueError(f"Missing required columns: {missing}")
        return record[FEATURE_ORDER]
    if isinstance(record, Mapping):
        missing = [name for name in FEATURE_ORDER if name not in record]
        if missing:
            raise ValueError(f"Missing required features: {missing}")
        return pd.DataFrame([[record[name] for name in FEATURE_ORDER]], columns=FEATURE_ORDER)
    values = list(record)
    if len(values) != len(FEATURE_ORDER):
        raise ValueError(f"Expected {len(FEATURE_ORDER)} values, received {len(values)}")
    return pd.DataFrame([values], columns=FEATURE_ORDER)


def predict(record: Mapping[str, Any] | Sequence[Any] | pd.DataFrame) -> dict[str, Any]:
    """Predict one Covertype class from the original 54-feature order."""

    model = joblib.load(ROOT / "model.joblib")
    features = _as_frame(record)
    class_id = int(model.predict(features)[0])
    probabilities = model.predict_proba(features)[0]
    return {
        "label": LABELS[class_id],
        "class_id": class_id,
        "probabilities": {
            str(int(label)): float(probability)
            for label, probability in zip(model.classes_, probabilities)
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--json", required=True, help="JSON object with all 54 named features")
    args = parser.parse_args()
    print(json.dumps(predict(json.loads(args.json)), indent=2))


if __name__ == "__main__":
    main()