#!/usr/bin/env python3
"""
Honest experiment: does rtl_433 IQ data help the category classifier?
=====================================================================

The statistical category classifier is starved on minority classes
(train counts ~ Doorbell 39, Weather 12, Security 2). Local Flipper .sub
corpora are byte-identical mirrors, so they add nothing. The rtl_433 project
ships genuinely-independent labelled OOK captures (``.cu8`` IQ), heavy on
exactly the starved classes (weather + PIR/door/smoke security).

This script measures — honestly — whether adding those demodulated captures to
TRAINING helps, without fooling ourselves:

  * Flipper captures are split group-disjoint into train/test (the real
    product distribution; prod input is Flipper .sub RAW).
  * rtl_433 captures are demodulated (``rtl433_iq_demod.demod_file``), split
    group-disjoint by device folder, and used to AUGMENT training only.
  * We compare, on the SAME held-out FLIPPER test set:
        baseline  = train on Flipper-only
        augmented = train on Flipper + rtl_433
    Regression check: do the well-populated Flipper classes
    (Garage/Remote/Fan/Doorbell) get WORSE? If augmentation tanks them, reject.
  * New-capability check: on a held-out rtl_433 test set, can the augmented
    model recognise Weather/Security at all (classes Flipper can't measure —
    too few held-out samples)?

Nothing here touches the serving path; it only decides whether to retrain.
"""

from __future__ import annotations

import random
import warnings
from collections import Counter
from pathlib import Path

import numpy as np

warnings.filterwarnings("ignore")

from scripts.rtl433_iq_demod import demod_file
from scripts.train_category_classifier import (
    collect as collect_flipper, DEFAULT_DATASET, FEATURE_NAMES, extract_features,
)
from src.matcher.timing_analyzer import get_timing_analyzer
from src.matcher.preamble_detector import get_preamble_detector

RTL433_ROOT = (Path(__file__).parent.parent
               / "data/rf_test_datasets/rtl_433_tests/tests")

# Conservative folder -> category map. Only clearly in-taxonomy devices; the
# demod quality-gate additionally drops any FSK/degenerate captures. TPMS,
# energy/utility meters, thermostats and unknowns are intentionally omitted.
RTL433_FOLDER_TO_CATEGORY = {
    # Weather Sensor (the biggest starved-class win)
    **{f: "Weather Sensor" for f in [
        "acurite", "alectov1", "alecto_ws_1200", "ambient_weather", "auriol",
        "bresser_3ch", "bresser_5in1", "bresser_6in1", "bt_rain", "calibeur",
        "companion_wtr001", "conrad_pool_thermometer", "cotech", "cresta_ws688",
        "ecowitt", "EcoWitt-WH40", "esperanza_ews", "Eurochron-EFTH800",
        "eurochron-th", "fineoffset", "froggit_wh1080_Pass14c", "ft004b",
        "hideki", "holman_ws5029", "imagintronix_wh5", "infactory_PV-8796-675",
        "InFactroy-Temp-Humidity-Sensor-T05K-THC", "inkbird", "inovalley-kw9015b",
        "lacrosse", "lacrosse_ltv", "lacrosse_ws7000", "maverick_et-73",
        "maverick_et733", "mebus_te204nl", "misol", "missil_ml0757", "nexus",
        "Opus-XT300", "oregon_scientific", "prologue", "proove", "rubicson",
        "s3318p", "sharp_spc344", "sharp_spc775", "solight_te44", "springfield",
        "TFA_30.3196_TempHumid", "tfa_30_3211_02", "TFA-Drop-30.3233.01",
        "TFA_Marbella", "TFA-Pool-thermometer-30.3160", "TFA-Twin-Plus-30.3049",
        "thermopro-tp11", "thermopro-tp12", "thermopro-tx2", "TS-FT002",
        "ttx201", "tx22-it", "wssensor", "wt0124", "XC-0324", "xc0348",
        "Globaltronics", "PoolWT0122", "rh787t",
    ]},
    # Security Sensor (motion / PIR / smoke / door / alarm)
    **{f: "Security Sensor" for f in [
        "generic_door_sensor", "generic_motion", "honeywell", "honeywell_5816",
        "honeywell_5890PI", "honeywell_activlink", "dsc", "simplisafe",
        "skylink_motion", "smoke_gs558", "cavius", "KIDDE_RF-SM-ACDC",
        "visonic_powercode", "interlogix", "x10_sec", "EV1527-PIR-SGOOWAY",
        "Kerui-D026", "chuango",
    ]},
    # Doorbell
    **{f: "Doorbell" for f in [
        "Byron-BY101", "Byron-BY34", "Byron_db304", "door_bell", "quhwa",
    ]},
    # Remote Control (remotes, outlet switches, key fobs)
    **{f: "Remote Control" for f in [
        "akhan_remote", "dish_remote_6.3", "directv", "EV1527-Universal-Remote",
        "fordremote", "generic-4ch-black-remote-315mhz",
        "generic-4ch-black-remote-433mhz", "generic_remote",
        "ge-coloreffects-remote", "honda_remote", "HT680_remote",
        "hyundai_remote", "intertechno", "nexa_LMST-606", "newkaku", "PT2262",
        "rayrun_rm03", "waveman", "x10", "philips", "Philips-AJ7010",
        "status-rc202", "kangtai", "kedsum", "etekcity", "elro", "blyss",
        "brennstuhl_rcs_2044", "telldus", "smarthome", "silvercrest",
        "uni-com-66125", "valeo",
    ]},
    # Garage Door Opener / gate
    **{f: "Garage Door Opener" for f in [
        "cardin", "LiftMaster_4330E", "secplus", "keeloq", "Microchip-HCS200",
    ]},
    # Fan Controller
    **{f: "Fan Controller" for f in [
        "lucci-air-ceiling-fan-remote", "SQMLtdFan",
    ]},
}


def collect_rtl433(seed: int, per_folder: int = 40):
    """Demodulate rtl_433 captures into (X, y, groups) using the folder map."""
    rng = random.Random(seed)
    ta, pd = get_timing_analyzer(), get_preamble_detector()
    X, y, groups = [], [], []
    skipped = Counter()
    for folder, category in RTL433_FOLDER_TO_CATEGORY.items():
        fpath = RTL433_ROOT / folder
        if not fpath.is_dir():
            skipped["missing_folder"] += 1
            continue
        cu8s = sorted(fpath.rglob("*.cu8"))
        rng.shuffle(cu8s)
        taken = 0
        for c in cu8s:
            if taken >= per_folder:
                break
            pulses, freq = demod_file(c)
            if pulses is None:
                skipped["rejected_demod"] += 1
                continue
            feats = extract_features(pulses, freq, ta, pd)
            if feats is None:
                skipped["no_timing"] += 1
                continue
            X.append(feats)
            y.append(category)
            groups.append(f"rtl433:{folder}")  # group per device folder
            taken += 1
    return np.array(X), np.array(y), np.array(groups), skipped


def per_class_recall(y_true, y_pred, labels):
    out = {}
    for lab in labels:
        idx = y_true == lab
        n = int(idx.sum())
        out[lab] = (int((y_pred[idx] == lab).sum()), n)
    return out


def main():
    from sklearn.model_selection import GroupShuffleSplit
    from sklearn.ensemble import GradientBoostingClassifier
    from sklearn.metrics import balanced_accuracy_score, accuracy_score

    seed = 42
    print("Collecting Flipper RAW ...")
    Xf, yf, gf, ff, _ = collect_flipper(DEFAULT_DATASET, 400, seed)
    print(f"  Flipper: {len(Xf)} samples  {dict(Counter(yf))}")

    print("Demodulating rtl_433 IQ ...")
    Xr, yr, gr, sk = collect_rtl433(seed)
    print(f"  rtl_433: {len(Xr)} samples  {dict(Counter(yr))}  (skipped {dict(sk)})")

    labels = sorted(set(yf) | set(yr))

    # Group-disjoint Flipper train/test (the real product distribution).
    gss = GroupShuffleSplit(n_splits=1, test_size=0.25, random_state=seed)
    ftr, fte = next(gss.split(Xf, yf, gf))
    Xf_tr, Xf_te = Xf[ftr], Xf[fte]
    yf_tr, yf_te = yf[ftr], yf[fte]

    # Group-disjoint rtl_433 train/test.
    rtr, rte = next(GroupShuffleSplit(1, test_size=0.30, random_state=seed)
                    .split(Xr, yr, gr))
    Xr_tr, Xr_te = Xr[rtr], Xr[rte]
    yr_tr, yr_te = yr[rtr], yr[rte]

    def fit(X, y):
        return GradientBoostingClassifier(random_state=seed).fit(X, y)

    base = fit(Xf_tr, yf_tr)
    aug = fit(np.vstack([Xf_tr, Xr_tr]), np.concatenate([yf_tr, yr_tr]))

    print("\n" + "=" * 70)
    print("A) REGRESSION CHECK — held-out FLIPPER test (does augmentation hurt?)")
    print("=" * 70)
    for name, model in [("baseline (Flipper-only)", base),
                        ("augmented (Flipper+rtl433)", aug)]:
        p = model.predict(Xf_te)
        print(f"\n{name}:  bal={balanced_accuracy_score(yf_te, p):.3f} "
              f"top1={accuracy_score(yf_te, p):.3f}")
        for lab, (ok, n) in per_class_recall(yf_te, p, labels).items():
            if n:
                print(f"    {lab:20} {ok:3d}/{n:<3d}")

    print("\n" + "=" * 70)
    print("B) NEW-CAPABILITY — held-out rtl_433 test (classes Flipper can't measure)")
    print("=" * 70)
    for name, model in [("baseline (Flipper-only)", base),
                        ("augmented (Flipper+rtl433)", aug)]:
        p = model.predict(Xr_te)
        print(f"\n{name}:  bal={balanced_accuracy_score(yr_te, p):.3f} "
              f"top1={accuracy_score(yr_te, p):.3f}")
        for lab, (ok, n) in per_class_recall(yr_te, p, labels).items():
            if n:
                print(f"    {lab:20} {ok:3d}/{n:<3d}")


if __name__ == "__main__":
    main()
