Compressor/backend/diagnose_drift.py

from __future__ import annotations

import os
import sys
from time import monotonic

import numpy as np
import pandas as pd
import pymysql

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from app.config import settings  # noqa: E402

BATCHES = (30, 31, 32)
COLUMN = "YSJ_1"
RNG = np.random.default_rng(7)


def connect() -> pymysql.connections.Connection:
    return pymysql.connect(
        host=settings.db_host,
        port=settings.db_port,
        user=settings.db_user,
        password=settings.db_password,
        database=settings.db_name,
        charset="utf8mb4",
        cursorclass=pymysql.cursors.SSCursor,
        connect_timeout=settings.db_connect_timeout,
        read_timeout=3600,
        write_timeout=120,
    )


LEVEL_CACHE = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "cache", "pks_levels.csv"
)


def wall_offset() -> int:
    """Return seconds to add to UNIX epoch so naive pandas shows local wall-clock times."""
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute("SELECT sample_time, UNIX_TIMESTAMP(sample_time) AS u "
                           "FROM pks_long_sample LIMIT 1")
            wall, uni = cursor.fetchone()
    finally:
        connection.close()
    wall_epoch = int(pd.Timestamp(wall).tz_localize(None).value // 10 ** 9)
    return wall_epoch - int(uni)


def load_levels(offset: int) -> pd.DataFrame:
    """1-min bucket AVERAGE level per batch, from a single server-side aggregation (cached)."""
    if os.path.exists(LEVEL_CACHE):
        frame = pd.read_csv(LEVEL_CACHE)
        frame["time"] = pd.to_datetime(frame["time"])
        print(f"   level buckets loaded from cache: {len(frame):,} rows")
        return frame
    t0 = monotonic()
    connection = connect()
    rows = []
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT import_batch_id AS b, UNIX_TIMESTAMP(sample_time) DIV 60 AS bk, "
                f"COUNT(*) AS cnt, AVG(`{COLUMN}`) AS lvl "
                "FROM pks_long_sample GROUP BY import_batch_id, bk ORDER BY import_batch_id, bk"
            )
            while True:
                chunk = cursor.fetchmany(200_000)
                if not chunk:
                    break
                rows.extend(chunk)
    finally:
        connection.close()
    frame = pd.DataFrame(rows, columns=["b", "bk", "cnt", "lvl"])
    frame["time"] = pd.to_datetime(frame["bk"] * 60 + offset, unit="s")
    bad = frame[frame["cnt"] != 60]
    if len(bad):
        print(f"   WARN: {len(bad)} buckets with cnt!=60\n{bad.head(10).to_string()}")
    frame[["b", "bk", "cnt", "lvl", "time"]].to_csv(LEVEL_CACHE, index=False)
    print(f"   level buckets loaded: {len(frame):,} rows ({frame['b'].nunique()} batches) "
          f"{monotonic()-t0:.0f}s")
    return frame


def sample_raw_days() -> pd.DataFrame:
    """Fetch raw 5s rows of ~24 random days per batch to gauge noise/quantization/gaps."""
    t0 = monotonic()
    days = pd.date_range("2025-04-01", "2026-04-30", freq="D")
    picks = days[RNG.choice(len(days), size=24, replace=False)].sort_values()
    connection = connect()
    parts = []
    try:
        with connection.cursor() as cursor:
            for d in picks:
                start = d.strftime("%Y-%m-%d 00:00:00")
                end = d.strftime("%Y-%m-%d 23:59:55")
                cursor.execute(
                    "SELECT import_batch_id AS b, "
                    f"UNIX_TIMESTAMP(sample_time) AS t, `{COLUMN}` AS v "
                    "FROM pks_long_sample WHERE sample_time BETWEEN %s AND %s",
                    (start, end),
                )
                while True:
                    chunk = cursor.fetchmany(100_000)
                    if not chunk:
                        break
                    parts.extend(chunk)
    finally:
        connection.close()
    frame = pd.DataFrame(parts, columns=["b", "t", "v"])
    frame = frame.dropna(subset=["v"]).sort_values(["b", "t"]).reset_index(drop=True)
    print(f"   raw sampled days: {len(frame):,} rows over {len(picks)} days "
          f"({monotonic()-t0:.0f}s)")
    return frame


def noise_stats(raw: pd.DataFrame) -> dict:
    out = {}
    for b, g in raw.groupby("b"):
        dv = np.abs(np.diff(g["v"].to_numpy()))
        pos = dv[dv > 1e-12]
        med = float(np.median(dv))
        mad0 = float(np.median(np.abs(dv - med)))
        out[int(b)] = {
            "n_pairs": int(len(dv)),
            "pct_eq0": round(float((dv <= 1e-12).mean()), 4),
            "adj_diff": {
                "median": med,
                "mad": mad0,
                "delta(med+3*mad)": round(med + 3 * mad0, 4),
                "min_pos": float(pos.min()) if pos.size else 0.0,
                "p1_pos": round(float(np.percentile(pos, 1)), 5) if pos.size else 0.0,
                "p50_pos": round(float(np.percentile(pos, 50)), 5) if pos.size else 0.0,
                "p99_pos": round(float(np.percentile(pos, 99)), 5) if pos.size else 0.0,
            },
            "value": {
                "min": round(float(g["v"].min()), 3),
                "p1": round(float(g["v"].quantile(0.01)), 3),
                "p50": round(float(g["v"].median()), 3),
                "p99": round(float(g["v"].quantile(0.99)), 3),
                "max": round(float(g["v"].max()), 3),
            },
        }
    return out


def series_per_batch(levels: pd.DataFrame) -> dict[int, pd.Series]:
    out = {}
    for b, g in levels.groupby("b"):
        s = g.set_index("time")["lvl"]
        s = s[~s.index.duplicated(keep="first")].sort_index()
        out[int(b)] = s
    return out


def run_stats(values: np.ndarray, tau: float, minutes: int = 5) -> dict:
    mask = np.abs(values) > tau
    if mask.size == 0:
        return {"count": 0, "frac": 0.0, "top": []}
    diff = np.diff(mask.astype(np.int8))
    starts = np.flatnonzero(diff == 1) + 1
    ends = np.flatnonzero(diff == -1) + 1
    if mask[0]:
        starts = np.concatenate(([0], starts))
    if mask[-1]:
        ends = np.concatenate((ends, [len(mask)]))
    if len(starts) == 0 or len(ends) == 0:
        return {"count": 0, "frac": float(mask.mean()), "top": []}
    ends = ends[: len(starts)]
    durs = (ends - starts) * minutes
    peaks = [float(np.max(np.abs(values[s:e]))) for s, e in zip(starts, ends)]
    order = np.argsort(peaks)[::-1][:6]
    top = [{"peak": round(peaks[i], 3), "dur_min": int(durs[i])} for i in order]
    return {
        "count": int(len(durs)),
        "durs_min": [int(x) for x in durs],
        "frac": round(float(mask.mean()), 4),
        "top": top,
    }


def main() -> None:
    out: dict = {}
    coverage = {}
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT import_batch_id AS b, COUNT(*) AS n, COUNT(YSJ_1) AS nn, "
                "MIN(sample_time) AS mn, MAX(sample_time) AS mx, "
                "COUNT(DISTINCT sample_time) AS dd FROM pks_long_sample GROUP BY import_batch_id"
            )
            for row in cursor.fetchall():
                b, n, nn, mn, mx, dd = row
                span = int((mx - mn).total_seconds() // 5) + 1
                coverage[int(b)] = {
                    "rows": int(n), "non_null": int(nn), "distinct_t": int(dd),
                    "span_5s_slots": span, "gap_free": n == dd == span,
                }
    finally:
        connection.close()
    out["coverage"] = coverage
    print("coverage:", out["coverage"], flush=True)

    offset = wall_offset()
    print(f"wall offset = {offset}s (UTC -> local)")
    levels = load_levels(offset)
    per_batch = series_per_batch(levels)

    raw = sample_raw_days()
    out["noise_and_quantization"] = noise_stats(raw)

    self_stats: dict[int, dict] = {}
    for b in BATCHES:
        s5 = per_batch[b].resample("5min").median().dropna()
        local6h = s5.rolling(145, center=True, min_periods=20).median()
        ref10d = s5.rolling(5761, center=True, min_periods=100).median()
        e = (local6h - ref10d).dropna().to_numpy()
        med_e = float(np.median(e))
        sigma_e = 1.4826 * float(np.median(np.abs(e - med_e)))
        self_stats[b] = {
            "med": round(med_e, 4), "sigma": round(sigma_e, 4), "tau4": round(4 * sigma_e, 4),
            "max_abs": round(float(np.abs(e).max()), 4),
            "flag": {f"{k}s": run_stats(e, k * sigma_e) for k in (3, 4, 5, 6)},
            "daily_level_p50": round(float(per_batch[b].resample("1D").median().median()), 4),
            "daily_level_minmax": [
                round(float(per_batch[b].resample("1D").median().min()), 3),
                round(float(per_batch[b].resample("1D").median().max()), 3),
            ],
        }
        print(f"self b{b}: med={self_stats[b]['med']} sigma={self_stats[b]['sigma']} "
              f"tau4={self_stats[b]['tau4']} max={self_stats[b]['max_abs']} "
              f"frac@4s={self_stats[b]['flag']['4s']['frac']} runs@4s={self_stats[b]['flag']['4s']['count']}")
    out["self_ref"] = self_stats

    s5_all = {b: per_batch[b].resample("5min").median() for b in BATCHES}
    frame = pd.DataFrame(s5_all).dropna()
    cross = {}
    for u in BATCHES:
        others = [o for o in BATCHES if o != u]
        ref = frame[others].median(axis=1)
        d = (frame[u] - ref).to_numpy()
        med_d = float(np.median(d))
        sigma_d = 1.4826 * float(np.median(np.abs(d - med_d)))
        cross[int(u)] = {
            "med": round(med_d, 4), "sigma": round(sigma_d, 4), "tau4": round(4 * sigma_d, 4),
            "max_abs": round(float(np.abs(d).max()), 4),
            "flag": {f"{k}s": run_stats(d, k * sigma_d) for k in (3, 4, 5, 6)},
        }
        print(f"cross b{u}: med={cross[u]['med']} sigma={cross[u]['sigma']} "
              f"tau4={cross[u]['tau4']} max={cross[u]['max_abs']} "
              f"frac@4s={cross[u]['flag']['4s']['frac']} runs@4s={cross[u]['flag']['4s']['count']}")
    out["cross_ref"] = cross
    out["cross_grid"] = {"common_5min": int(len(frame)),
                         "start": str(frame.index.min()), "end": str(frame.index.max())}

    print("\n===== JSON SUMMARY =====")
    print(__import__("json").dumps(out, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()