main.py

"""CLI entry point for the PCS predictor pipeline."""
import argparse
import logging
import sys
from pathlib import Path

from db.database import init_db, get_conn, upsert_race, upsert_stage, upsert_rider, insert_result
from scraper.races import iter_races, fetch_race_stages
from scraper.results import fetch_stage_results
from scraper.riders import fetch_rider
from config import SCRAPE_YEARS, DB_PATH, WOMEN_RACE_CLASSES, WOMEN_CIRCUIT, MEN_CIRCUITS

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)


def cmd_init(_args):
    """Initialise the database."""
    init_db()
    log.info("database initialised at %s", DB_PATH)


def cmd_scrape(args):
    """Scrape races, stages, results, and rider profiles."""
    init_db()
    years = [int(y) for y in args.years.split(",")] if args.years else SCRAPE_YEARS
    if args.gender == "women":
        circuits = {WOMEN_CIRCUIT: WOMEN_RACE_CLASSES}
    else:
        circuits = MEN_CIRCUITS

    log.info("scraping races for years: %s (gender=%s)", years, args.gender)
    for race in iter_races(years, circuits=circuits):
        # Commit per-race so progress is never lost on crash/interrupt
        with get_conn() as conn:
            race_id = upsert_race(conn, race)
            log.info("race: %s (id=%d)", race["name"], race_id)

            stages = fetch_race_stages(race["pcs_slug"])
            for stage_data in stages:
                if not stage_data.get("pcs_slug"):
                    log.warning("skipping stage with empty slug in %s", race["name"])
                    continue
                stage_data["race_id"] = race_id
                stage_id = upsert_stage(conn, stage_data)

                results = fetch_stage_results(stage_data["pcs_slug"])
                riders_seen: dict[str, int] = {}

                for r in results:
                    slug = r["rider_slug"]
                    if slug not in riders_seen:
                        if not args.skip_riders:
                            rider = fetch_rider(slug)
                            if rider:
                                rider_id = upsert_rider(conn, rider)
                            else:
                                rider_id = upsert_rider(conn, {
                                    "pcs_slug": slug, "name": r["rider_name"],
                                    "nationality": None, "dob": None, "team": None,
                                    "pcs_rank": None, "speciality": None,
                                    "weight_kg": None, "height_cm": None,
                                })
                        else:
                            rider_id = upsert_rider(conn, {
                                "pcs_slug": slug, "name": r["rider_name"],
                                "nationality": None, "dob": None, "team": None,
                                "pcs_rank": None, "speciality": None,
                                "weight_kg": None, "height_cm": None,
                            })
                        riders_seen[slug] = rider_id

                    insert_result(conn, {
                        "stage_id": stage_id,
                        "rider_id": riders_seen[slug],
                        "position": r["position"],
                        "status": r["status"],
                        "time_seconds": r["time_seconds"],
                        "points_pcs": r["points_pcs"],
                        "points_uci": r["points_uci"],
                        "bib": r["bib"],
                    })

    log.info("scraping complete")


def cmd_train(args):
    """Train the top-10 probability model."""
    from model.train import train
    train(train_cutoff=args.cutoff, val_race_slug=args.val_race, gender=args.gender)


def cmd_backtest(args):
    """Backtest predictions against a historic race."""
    from model.backtest import backtest
    backtest(args.race, args.cutoff, args.top_n, gender=args.gender)


def cmd_predict(args):
    """Predict top-10 probabilities for a race."""
    from model.predict import predict_race, print_predictions
    df = predict_race(args.race, args.cutoff, gender=args.gender)
    if df.empty:
        log.error("no predictions generated")
        return
    print_predictions(df, top_n=args.top_n)


def cmd_scrape_elevation(args):
    """Backfill elevation_m for stages that are missing it.

    Also fixes NULL date/distance/profile_type for one-day stages whose metadata
    wasn't scraped correctly (e.g. because _parse_infolist used div.bold instead
    of div.title).
    """
    from scraper.races import fetch_stage_elevation, fetch_race_stages, fetch_result_meta
    from db.database import get_conn

    # ── fix one-day stages with NULL date (broken metadata parse) ─────────────
    with get_conn() as conn:
        null_meta = conn.execute(
            """SELECT s.id, s.pcs_slug, r.is_stage_race
               FROM stages s JOIN races r ON s.race_id = r.id
               WHERE s.date IS NULL
               ORDER BY s.id"""
        ).fetchall()

    log.info("stages with NULL date: %d", len(null_meta))
    meta_fixed = 0
    for row in null_meta:
        stages = fetch_race_stages(row["pcs_slug"]) if not row["is_stage_race"] else []
        if stages and stages[0].get("date"):
            s = stages[0]
            with get_conn() as conn:
                conn.execute(
                    """UPDATE stages SET date=?, distance_km=?, elevation_m=?,
                       profile_type=?, departure=?, arrival=?
                       WHERE id=?""",
                    (s["date"], s["distance_km"], s["elevation_m"],
                     s["profile_type"], s["departure"], s["arrival"], row["id"]),
                )
            meta_fixed += 1
            log.info("  meta fixed: %s → %s", row["pcs_slug"], s["date"])
        else:
            log.debug("  no meta found: %s", row["pcs_slug"])

    log.info("meta backfill complete: %d/%d fixed", meta_fixed, len(null_meta))

    # ── backfill elevation_m ───────────────────────────────────────────────────
    with get_conn() as conn:
        rows = conn.execute(
            """SELECT id, pcs_slug FROM stages
               WHERE elevation_m IS NULL AND pcs_slug NOT LIKE '%/result'
               ORDER BY id"""
        ).fetchall()

    log.info("stages missing elevation: %d", len(rows))
    updated = 0
    for row in rows:
        elev = fetch_stage_elevation(row["pcs_slug"])
        if elev:
            with get_conn() as conn:
                conn.execute(
                    "UPDATE stages SET elevation_m = ? WHERE id = ?",
                    (elev, row["id"]),
                )
            updated += 1
            log.info("  %s → %d m", row["pcs_slug"], elev)
        else:
            log.debug("  %s → no data", row["pcs_slug"])

    log.info("elevation backfill complete: %d/%d updated", updated, len(rows))

    # ── backfill gradient_final_km / profile_score / profile_type ─────────────
    with get_conn() as conn:
        grad_rows = conn.execute(
            """SELECT id, pcs_slug FROM stages
               WHERE gradient_final_km IS NULL
               ORDER BY id"""
        ).fetchall()

    log.info("stages missing gradient/profile meta: %d", len(grad_rows))
    meta_updated = 0
    for row in grad_rows:
        meta = fetch_result_meta(row["pcs_slug"])
        if meta.get("gradient_final_km") is not None or meta.get("profile_score") is not None:
            with get_conn() as conn:
                conn.execute(
                    """UPDATE stages SET
                       gradient_final_km = COALESCE(?, gradient_final_km),
                       profile_score     = COALESCE(?, profile_score),
                       profile_type      = COALESCE(profile_type, ?)
                       WHERE id = ?""",
                    (meta.get("gradient_final_km"), meta.get("profile_score"),
                     meta.get("profile_type"), row["id"]),
                )
            meta_updated += 1
            log.info("  %s → grad=%.1f%% score=%s profile=%s",
                     row["pcs_slug"],
                     meta.get("gradient_final_km") or 0,
                     meta.get("profile_score"),
                     meta.get("profile_type"))
        else:
            log.debug("  no meta: %s", row["pcs_slug"])

    log.info("gradient/profile backfill: %d/%d updated", meta_updated, len(grad_rows))


def cmd_scrape_race(args):
    """Scrape stages, results, and optionally rider profiles for a single race slug."""
    init_db()
    race_slug = args.race_slug

    from scraper.races import fetch_race_stages

    with get_conn() as conn:
        row = conn.execute(
            "SELECT id FROM races WHERE pcs_slug = ?", (race_slug,)
        ).fetchone()

    if not row:
        log.error(
            "Race '%s' not found in DB. Run 'scrape' first to populate the calendar, "
            "or check the slug.",
            race_slug,
        )
        return

    race_id = row["id"]
    stages = fetch_race_stages(race_slug)
    if not stages:
        log.error("No stages found for %s", race_slug)
        return

    log.info("scraping %d stage(s) for %s", len(stages), race_slug)
    total_results = 0

    for stage_data in stages:
        if not stage_data.get("pcs_slug"):
            log.warning("skipping stage with empty slug in %s", race_slug)
            continue
        stage_data["race_id"] = race_id
        with get_conn() as conn:
            stage_id = upsert_stage(conn, stage_data)

        results = fetch_stage_results(stage_data["pcs_slug"])
        if not results:
            log.warning("no results for stage %s", stage_data["pcs_slug"])
            continue

        riders_seen: dict[str, int] = {}
        with get_conn() as conn:
            for r in results:
                slug = r["rider_slug"]
                if slug not in riders_seen:
                    if not args.skip_riders:
                        rider = fetch_rider(slug)
                        if rider:
                            rider_id = upsert_rider(conn, rider)
                        else:
                            rider_id = upsert_rider(conn, {
                                "pcs_slug": slug, "name": r["rider_name"],
                                "nationality": None, "dob": None, "team": None,
                                "pcs_rank": None, "speciality": None,
                                "weight_kg": None, "height_cm": None,
                            })
                    else:
                        rider_id = upsert_rider(conn, {
                            "pcs_slug": slug, "name": r["rider_name"],
                            "nationality": None, "dob": None, "team": None,
                            "pcs_rank": None, "speciality": None,
                            "weight_kg": None, "height_cm": None,
                        })
                    riders_seen[slug] = rider_id

                insert_result(conn, {
                    "stage_id": stage_id,
                    "rider_id": riders_seen[slug],
                    "position": r["position"],
                    "status": r["status"],
                    "time_seconds": r["time_seconds"],
                    "points_pcs": r["points_pcs"],
                    "points_uci": r["points_uci"],
                    "bib": r["bib"],
                })
            total_results += len(results)
            log.info("  stage %s → %d results", stage_data["pcs_slug"], len(results))

    log.info("done: %d total results for %s", total_results, race_slug)


def cmd_tag_surface(args):
    """Backfill surface column for existing stages based on race slug."""
    from db.database import get_conn
    from config import COBBLED_RACE_SLUGS, GRAVEL_RACE_SLUGS

    with get_conn() as conn:
        stages = conn.execute(
            "SELECT s.id, s.pcs_slug FROM stages s"
        ).fetchall()

        cobbled = updated = 0
        for row in stages:
            slug = row["pcs_slug"]
            surface = "road"
            if any(p in slug for p in COBBLED_RACE_SLUGS):
                surface = "cobbled"
            elif any(p in slug for p in GRAVEL_RACE_SLUGS):
                surface = "gravel"
            if surface != "road":
                conn.execute(
                    "UPDATE stages SET surface = ? WHERE id = ?",
                    (surface, row["id"]),
                )
                cobbled += 1
            updated += 1

    log.info("surface tagging complete: %d/%d non-road stages", cobbled, updated)


def cmd_ingest_gpx(args):
    """Parse a GPX file and attach climb data to a stage."""
    from scraper.gpx import parse_gpx
    from db.database import get_conn

    gpx_path = Path(args.gpx)
    if not gpx_path.exists():
        log.error("file not found: %s", gpx_path)
        sys.exit(1)

    data = parse_gpx(gpx_path)
    log.info(
        "GPX summary — distance: %.1f km, gain: %.0f m, climbs: %d",
        data.get("total_distance_km", 0),
        data.get("total_elevation_gain_m", 0),
        len(data.get("climbs", [])),
    )

    if args.stage_slug:
        with get_conn() as conn:
            row = conn.execute(
                "SELECT id FROM stages WHERE pcs_slug = ?", (args.stage_slug,)
            ).fetchone()
            if not row:
                row = conn.execute(
                    "SELECT id FROM upcoming_stages WHERE pcs_slug = ?", (args.stage_slug,)
                ).fetchone()
            if not row:
                log.error("stage slug not found in DB: %s", args.stage_slug)
                sys.exit(1)

            stage_id = row["id"]
            # Store path reference
            conn.execute(
                "UPDATE stages SET gpx_path = ? WHERE id = ?", (str(gpx_path), stage_id)
            )
            # Insert climb records
            conn.execute("DELETE FROM stage_climbs WHERE stage_id = ?", (stage_id,))
            for i, climb in enumerate(data.get("climbs", []), start=1):
                conn.execute(
                    """INSERT INTO stage_climbs
                       (stage_id, climb_order, start_km, length_km,
                        elevation_gain_m, avg_gradient_pct, max_gradient_pct)
                       VALUES (?, ?, ?, ?, ?, ?, ?)""",
                    (stage_id, i, climb["start_km"], climb["length_km"],
                     climb["elevation_gain_m"], climb["avg_gradient_pct"],
                     climb["max_gradient_pct"]),
                )
        log.info("climbs stored for stage %s", args.stage_slug)


def main():
    parser = argparse.ArgumentParser(description="PCS Race Predictor")
    sub = parser.add_subparsers(dest="cmd", required=True)

    sub.add_parser("init", help="Initialise the database").set_defaults(func=cmd_init)

    p_scrape = sub.add_parser("scrape", help="Scrape PCS data")
    p_scrape.add_argument("--years", help="Comma-separated years, e.g. 2024,2025")
    p_scrape.add_argument("--skip-riders", action="store_true",
                          help="Don't fetch individual rider profiles (faster)")
    p_scrape.add_argument("--gender", default="men", choices=["men", "women"],
                          help="Which calendar to scrape (default: men)")
    p_scrape.set_defaults(func=cmd_scrape)

    p_train = sub.add_parser("train", help="Train the prediction model")
    p_train.add_argument("--cutoff", default="2024-05-04",
                         help="Train on data before this date (default: Giro 2024 start)")
    p_train.add_argument("--val-race", default="race/giro-d-italia/2024")
    p_train.add_argument("--gender", default="men", choices=["men", "women"])
    p_train.set_defaults(func=cmd_train)

    p_bt = sub.add_parser("backtest", help="Backtest against a historic race")
    p_bt.add_argument("--race",   default="race/giro-d-italia/2024")
    p_bt.add_argument("--cutoff", default="2024-05-04")
    p_bt.add_argument("--top-n",  type=int, default=10, dest="top_n")
    p_bt.add_argument("--gender", default="men", choices=["men", "women"])
    p_bt.set_defaults(func=cmd_backtest)

    p_pred = sub.add_parser("predict", help="Predict top-10 for an upcoming race")
    p_pred.add_argument("race",   help="Race slug, e.g. race/tour-de-france/2026")
    p_pred.add_argument("--cutoff", required=True, help="Feature cutoff date (race start)")
    p_pred.add_argument("--top-n", type=int, default=15, dest="top_n")
    p_pred.add_argument("--gender", default="men", choices=["men", "women"])
    p_pred.set_defaults(func=cmd_predict)

    p_gpx = sub.add_parser("gpx", help="Ingest a GPX file")
    p_gpx.add_argument("gpx", help="Path to .gpx file")
    p_gpx.add_argument("--stage", dest="stage_slug",
                       help="PCS stage slug to attach climb data to")
    p_gpx.set_defaults(func=cmd_ingest_gpx)

    sub.add_parser("scrape-elevation", help="Backfill elevation_m for all stages") \
       .set_defaults(func=cmd_scrape_elevation)

    sub.add_parser("tag-surface", help="Backfill cobbled/gravel surface tags") \
       .set_defaults(func=cmd_tag_surface)

    p_scrape_race = sub.add_parser(
        "scrape-race",
        help="Scrape stages + results for a single race slug (race must already be in DB)",
    )
    p_scrape_race.add_argument("race_slug", help="PCS race slug, e.g. race/milano-sanremo/2026")
    p_scrape_race.add_argument("--skip-riders", action="store_true",
                               help="Don't fetch individual rider profiles (faster)")
    p_scrape_race.set_defaults(func=cmd_scrape_race)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()