PolyWeather/src/analysis/probability_snapshot_archive.py

from __future__ import annotations

import json
import os
from datetime import datetime
from typing import Any, Dict, List, Optional

from src.database.runtime_state import (
    ProbabilitySnapshotRepository,
    STATE_STORAGE_DUAL,
    STATE_STORAGE_SQLITE,
    get_state_storage_mode,
)

DEDUP_SCAN_LINES = 200
MU_THRESHOLD = 0.2
SIGMA_THRESHOLD = 0.15
MAX_SO_FAR_THRESHOLD = 0.2
_snapshot_repo = ProbabilitySnapshotRepository()


def _sf(value: Any) -> Optional[float]:
    if value is None:
        return None
    try:
        return float(value)
    except Exception:
        return None


def _compact_snapshot(distribution: Optional[List[Dict[str, Any]]]) -> List[Dict[str, Any]]:
    compact: List[Dict[str, Any]] = []
    for row in distribution or []:
        if not isinstance(row, dict):
            continue
        value = row.get("value")
        probability = row.get("probability")
        if value is None or probability is None:
            continue
        try:
            compact.append(
                {
                    "v": int(value),
                    "p": round(float(probability), 3),
                }
            )
        except Exception:
            continue
        if len(compact) >= 4:
            break
    return compact


def _top_bucket(snapshot: Optional[List[Dict[str, Any]]]) -> Optional[int]:
    best_value = None
    best_prob = -1.0
    for row in snapshot or []:
        if not isinstance(row, dict):
            continue
        value = row.get("v")
        prob = _sf(row.get("p"))
        if value is None or prob is None:
            continue
        if prob > best_prob:
            best_value = int(value)
            best_prob = prob
    return best_value


def _load_recent_rows(path: str, max_lines: int = DEDUP_SCAN_LINES) -> List[Dict[str, Any]]:
    if not os.path.exists(path):
        return []
    with open(path, "r", encoding="utf-8") as fh:
        lines = fh.readlines()[-max_lines:]
    rows = []
    for line in lines:
        line = line.strip()
        if not line:
            continue
        try:
            row = json.loads(line)
        except Exception:
            continue
        if isinstance(row, dict):
            rows.append(row)
    return rows


def load_snapshot_rows_for_day(
    city_name: str,
    target_date: str,
    archive_path: Optional[str] = None,
) -> List[Dict[str, Any]]:
    city_key = str(city_name or "").strip().lower()
    date_key = str(target_date or "").strip()
    if not city_key or not date_key:
        return []

    mode = get_state_storage_mode()
    if mode == STATE_STORAGE_SQLITE:
        return _snapshot_repo.load_rows_by_city_date(city_key, date_key)

    root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    path = archive_path or os.path.join(root_dir, "data", "probability_training_snapshots.jsonl")
    if not os.path.exists(path):
        if mode == STATE_STORAGE_DUAL:
            return _snapshot_repo.load_rows_by_city_date(city_key, date_key)
        return []

    rows: List[Dict[str, Any]] = []
    try:
        with open(path, "r", encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                try:
                    row = json.loads(line)
                except Exception:
                    continue
                if not isinstance(row, dict):
                    continue
                if str(row.get("city") or "").strip().lower() != city_key:
                    continue
                if str(row.get("date") or "").strip() != date_key:
                    continue
                rows.append(row)
    except Exception:
        return []

    rows.sort(key=lambda row: str(row.get("timestamp") or ""))
    return rows


def _should_skip_append(path: str, payload: Dict[str, Any]) -> bool:
    mode = get_state_storage_mode()
    if mode == STATE_STORAGE_SQLITE:
        recent_rows = _snapshot_repo.load_recent_rows(
            str(payload.get("city") or ""),
            str(payload.get("date") or ""),
            DEDUP_SCAN_LINES,
        )
    else:
        recent_rows = _load_recent_rows(path)
    city = payload.get("city")
    date_str = payload.get("date")
    if not city or not date_str:
        return False

    for row in reversed(recent_rows):
        if row.get("city") != city or row.get("date") != date_str:
            continue
        if row.get("peak_status") != payload.get("peak_status"):
            return False
        if row.get("probability_mode") != payload.get("probability_mode"):
            return False

        current_top = _top_bucket(payload.get("prob_snapshot"))
        previous_top = _top_bucket(row.get("prob_snapshot"))
        current_shadow_top = _top_bucket(payload.get("shadow_prob_snapshot"))
        previous_shadow_top = _top_bucket(row.get("shadow_prob_snapshot"))
        if current_top != previous_top or current_shadow_top != previous_shadow_top:
            return False

        if abs((_sf(payload.get("raw_mu")) or 0.0) - (_sf(row.get("raw_mu")) or 0.0)) > MU_THRESHOLD:
            return False
        if abs((_sf(payload.get("raw_sigma")) or 0.0) - (_sf(row.get("raw_sigma")) or 0.0)) > SIGMA_THRESHOLD:
            return False
        if abs((_sf(payload.get("max_so_far")) or 0.0) - (_sf(row.get("max_so_far")) or 0.0)) > MAX_SO_FAR_THRESHOLD:
            return False

        return True

    return False


def append_probability_snapshot(
    city_name: str,
    *,
    local_date: str,
    observation_time: Optional[str],
    temp_symbol: str,
    raw_mu: Optional[float],
    raw_sigma: Optional[float],
    deb_prediction: Optional[float],
    ens_data: Optional[Dict[str, Any]],
    current_forecasts: Optional[Dict[str, Any]],
    max_so_far: Optional[float],
    current_temp: Optional[float] = None,
    humidity: Optional[float] = None,
    wind_speed_kt: Optional[float] = None,
    visibility_mi: Optional[float] = None,
    local_hour: Optional[float] = None,
    peak_status: Optional[str],
    probabilities: Optional[List[Dict[str, Any]]],
    shadow_probabilities: Optional[List[Dict[str, Any]]],
    calibration_summary: Optional[Dict[str, Any]],
    archive_path: Optional[str] = None,
) -> None:
    city_key = str(city_name or "").strip().lower()
    if not city_key:
        return

    root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    path = archive_path or os.path.join(
        root_dir,
        "data",
        "probability_training_snapshots.jsonl",
    )

    calibration_summary = calibration_summary or {}
    ens_data = ens_data or {}
    current_forecasts = current_forecasts or {}
    timestamp = str(observation_time or datetime.utcnow().isoformat() + "Z").strip()

    payload = {
        "city": city_key,
        "timestamp": timestamp,
        "date": local_date,
        "temp_symbol": temp_symbol,
        "raw_mu": _sf(raw_mu),
        "raw_sigma": _sf(raw_sigma),
        "deb_prediction": _sf(deb_prediction),
        "ensemble": {
            "p10": _sf(ens_data.get("p10")),
            "median": _sf(ens_data.get("median")),
            "p90": _sf(ens_data.get("p90")),
        },
        "multi_model": {
            key: _sf(value)
            for key, value in current_forecasts.items()
            if _sf(value) is not None
        },
        "max_so_far": _sf(max_so_far),
        "observation": {
            "current_temp": _sf(current_temp),
            "humidity": _sf(humidity),
            "wind_speed_kt": _sf(wind_speed_kt),
            "visibility_mi": _sf(visibility_mi),
            "local_hour": _sf(local_hour),
        },
        "peak_status": peak_status,
        "prob_snapshot": _compact_snapshot(probabilities),
        "shadow_prob_snapshot": _compact_snapshot(shadow_probabilities),
        "probability_engine": calibration_summary.get("engine"),
        "probability_mode": calibration_summary.get("mode"),
        "calibration_version": calibration_summary.get("calibration_version"),
        "calibration_source": calibration_summary.get("calibration_source"),
        "calibrated_mu": _sf(calibration_summary.get("calibrated_mu")),
        "calibrated_sigma": _sf(calibration_summary.get("calibrated_sigma")),
    }

    parent = os.path.dirname(os.path.abspath(path))
    if parent:
        os.makedirs(parent, exist_ok=True)

    if _should_skip_append(path, payload):
        return

    mode = get_state_storage_mode()
    if mode in {STATE_STORAGE_DUAL, STATE_STORAGE_SQLITE}:
        _snapshot_repo.append_snapshot(payload)

    if mode != STATE_STORAGE_SQLITE:
        with open(path, "a", encoding="utf-8") as fh:
            fh.write(json.dumps(payload, ensure_ascii=False) + "\n")