data_analysis_os.py

"""
OpenSkill (Weng-Lin Bayesian) rating models.

Modelling approach
------------------
Every round is treated as a 2-team match (marines vs aliens) and rated with
one or more internal models from `openskill` (for example Plackett-Luce,
Thurstone-Mosteller, Bradley-Terry). Each player has a Gaussian belief about
their skill,
`Rating(mu, sigma)`: `mu` is the estimated skill and `sigma` the model's
uncertainty about that estimate. After every round, `model.rate()` shifts the
`mu` of the winning team's players up and the losing team's down (by an
amount proportional to how surprising the result was) and shrinks `sigma`
towards the players who took part.

Just like the AdaGrad model in data_analysis_mlt.py, this is inherently
sequential -- each round's update depends on the ratings produced by every
earlier round -- so it is run as a single pass over `filtered_round_ids` in
plain Python, not as a Spark job. Round data is pulled once via `.toPandas()`
up front (see the comment in data_analysis_mlt.py for why this matters for
performance).

For comparability with the other models, win probability for each round is
predicted with `model.predict_win()` *before* that round's rating update is
applied (a fair, online/no-lookahead evaluation), and the same
accuracy/ROC-AUC/confusion-matrix metrics are computed and persisted to the
shared `metrics` table (name="os" for legacy Plackett-Luce, `os_<variant>`
for additional variants).

Final per-player skill is each model's OpenSkill `ordinal()` value (mu - 3*sigma): a
conservative skill estimate for which the player's true skill is ~99.7%
likely to be higher, so newer/less-played players aren't rated above
well-established, consistently good players purely because of overconfidence.
"""

import argparse, os
from pyspark.sql import Row
from pyspark.sql.functions import col
from sklearn.metrics import confusion_matrix, log_loss, roc_auc_score
import openskill.models as openskill_models
from ea_common import spark, filtered_rounds, filtered_round_ids, logger, users, db_path
from ea_schema import metrics_schema, skill_os_schema

# Models to evaluate. `key` drives metric/table names.
# `pl` keeps legacy names (`os`, `skill_os`) for backward compatibility.
MODEL_CANDIDATES = [
    ("pl", "PlackettLuce"),
    ("tmf", "ThurstoneMostellerFull"),
    ("btf", "BradleyTerryFull"),
    ("btp", "BradleyTerryPart"),
]


def available_model_classes():
    available = {}
    for key, class_name in MODEL_CANDIDATES:
        model_class = getattr(openskill_models, class_name, None)
        if model_class is not None:
            available[key] = (class_name, model_class)
    return available


available = available_model_classes()
default_models = [key for key, _class_name in MODEL_CANDIDATES if key in available]

parser = argparse.ArgumentParser(description="OpenSkill rating models")
parser.add_argument("--limit", type=int, help="Limit number of processed rounds", default=-1)
parser.add_argument(
    "--model",
    action="append",
    dest="models",
    choices=sorted(available.keys()),
    help="Model key to run (repeatable). Defaults to all available: {}".format(
        ", ".join(default_models)
    ),
)
args = parser.parse_args()

round_ids_to_process = filtered_round_ids
if int(args.limit) > 0:
    round_ids_to_process = round_ids_to_process[:int(args.limit)]

selected_model_keys = args.models if args.models else default_models
if not selected_model_keys:
    raise RuntimeError("No supported OpenSkill models available in this openskill installation.")

# ---------------------------------------------------------------------------
# Pull just the columns needed for the sequential loop ONCE (see
# data_analysis_mlt.py for why running a Spark job per round would be far too
# slow), then run the round-by-round update in plain Python.
# ---------------------------------------------------------------------------
rounds_pd = filtered_rounds.select("round_id", "user_id", "team", "result").toPandas()
rounds_by_id = {
    round_id: list(group[["user_id", "team", "result"]].itertuples(index=False, name=None))
    for round_id, group in rounds_pd.groupby("round_id")
}

all_user_ids = [row.id for row in users.select("id").collect()]


def evaluate_model(model_key, model_class_name, model_class):
    model = model_class()
    ratings = {}  # user_id -> OpenSkill rating object

    def get_rating(user_id):
        if user_id not in ratings:
            ratings[user_id] = model.rating(name=str(user_id))
        return ratings[user_id]

    predicted_probs = []   # P(marines win), predicted BEFORE the round's update
    actual_results = []    # 1 if marines won, else 0

    n = 1
    for round_id in round_ids_to_process:
        logger.info("[{}] Processing round: {}".format(model_key, n))
        n += 1

        players = rounds_by_id.get(round_id, [])
        if not players:
            continue

        marine_ids = [user_id for user_id, team, _result in players if team == 1]
        alien_ids = [user_id for user_id, team, _result in players if team == -1]
        if not marine_ids or not alien_ids:
            continue

        result = players[0][2]  # same result value on every row of the round

        marine_team = [get_rating(uid) for uid in marine_ids]
        alien_team = [get_rating(uid) for uid in alien_ids]

        # Predict BEFORE updating so evaluation never peeks at the outcome the
        # update itself is derived from (fair, online evaluation).
        win_probs = model.predict_win([marine_team, alien_team])
        marine_win_prob = float(win_probs[0])
        predicted_probs.append(marine_win_prob)
        actual_results.append(result)

        # Lower rank = winner. result == 1 means marines (team 1) won.
        ranks = [0, 1] if result == 1 else [1, 0]
        [new_marine_team, new_alien_team] = model.rate([marine_team, alien_team], ranks=ranks)

        for uid, rating in zip(marine_ids, new_marine_team):
            ratings[uid] = rating
        for uid, rating in zip(alien_ids, new_alien_team):
            ratings[uid] = rating

        logger.debug(
            "[{}] round {}: p(marines win)={}, result={}".format(
                model_key, round_id, marine_win_prob, result,
            )
        )

    default_rating = model.rating()

    # Keep legacy names for PlackettLuce output, add suffixed names for other variants.
    metric_name = "os" if model_key == "pl" else "os_{}".format(model_key)
    table_name = "skill_os" if model_key == "pl" else "skill_os_{}".format(model_key)
    mu_col = "mu_os" if model_key == "pl" else "mu_os_{}".format(model_key)
    sigma_col = "sigma_os" if model_key == "pl" else "sigma_os_{}".format(model_key)
    skill_col = "skill_os" if model_key == "pl" else "skill_os_{}".format(model_key)

    schema = skill_os_schema(mu_col, sigma_col, skill_col)
    skill_rows = [
        (
            int(uid),
            float(ratings.get(uid, default_rating).mu),
            float(ratings.get(uid, default_rating).sigma),
            float(ratings.get(uid, default_rating).ordinal()),
        )
        for uid in all_user_ids
    ]
    skill_df = spark.createDataFrame(skill_rows, schema)

    # Show results ordered by new skill, joined with nicknames for readability
    skill_df \
        .join(users.select(col("id").alias("user_id"), "nickname"), "user_id", "left") \
        .orderBy(col(skill_col).desc()) \
        .show()

    # Save per-model skill table so plots can load each variant separately.
    skill_df.write.mode("overwrite").parquet(db_path(table_name))

    if not actual_results:
        logger.warning("[{}] No valid rounds to evaluate; skipping metrics write.".format(model_key))
        return

    predicted_results = [1 if p > 0.5 else 0 for p in predicted_probs]
    accuracy = sum(1 for pred, actual in zip(predicted_results, actual_results) if pred == actual) / len(actual_results)
    roc_auc = roc_auc_score(actual_results, predicted_probs)
    model_log_loss = log_loss(actual_results, predicted_probs, labels=[0, 1])
    true_neg, false_pos, false_neg, true_pos = confusion_matrix(
        actual_results, predicted_results, labels=[0, 1],
    ).ravel()

    milestone = len(filtered_round_ids)
    model_metrics = spark.createDataFrame([
        Row(name=metric_name, metric="accuracy", milestone=milestone, value=float(accuracy)),
        Row(name=metric_name, metric="roc_auc", milestone=milestone, value=float(roc_auc)),
        Row(name=metric_name, metric="log_loss", milestone=milestone, value=float(model_log_loss)),
        Row(name=metric_name, metric="true_positives", milestone=milestone, value=float(true_pos)),
        Row(name=metric_name, metric="false_positives", milestone=milestone, value=float(false_pos)),
        Row(name=metric_name, metric="true_negatives", milestone=milestone, value=float(true_neg)),
        Row(name=metric_name, metric="false_negatives", milestone=milestone, value=float(false_neg)),
    ], metrics_schema)

    if os.path.exists(db_path("metrics")):
        model_metrics.write.mode("append").parquet(db_path("metrics"))
    else:
        model_metrics.write.parquet(db_path("metrics"))

    model_metrics.show()
    print(
        "{} [{}] -> Accuracy: {}, ROC-AUC: {}, Log Loss: {}".format(
            model_class_name, metric_name, accuracy, roc_auc, model_log_loss,
        )
    )

for model_key in selected_model_keys:
    class_name, model_class = available[model_key]
    logger.info("Running OpenSkill model {} ({})".format(model_key, class_name))
    evaluate_model(model_key, class_name, model_class)

spark.stop()