ohmu/dev/varvusensor/collector/calibrate.py

"""Kalibreerimise abivahend AS7341 värviandurile.

Loeb ClickHouse'ist (vaikimisi uus67.varvus) ja toetab:
  * capture-dark:  salvesta tume referents (andur kaetud) -> calibration.json
  * capture-white: salvesta valge referents (valge pind, SAMA valgustus) -> calibration.json
  * show:          näita viimaseid mõõtmisi (toores vs tume+valge kalibreeritud)

Kalibreeritud reflekteeruvus:  R = (S - D) / (W - D)  iga kanali kohta.
  * Tume (D) eemaldab tumedavoolu ja plaadi toite-LED-i (roheline) panuse.
  * Valge (W) eemaldab valgustuse spektri ja kanalite tundlikkuse erinevuse.

Protseduur:
  1. Fikseeri valgustus ja geomeetria (nurk, kaugus) kogu seeria jooksul.
  2. Kata andur (täiesti pime) -> capture-dark --last 5
  3. Mõõda valge pind (nt valge paber/PTFE) -> capture-white --last 5
  4. Mõõda proovid -> show --last 10

Käivitamine (projekti .venv):
  CH_URL=http://100.87.1.24:8125 CH_USER=zigbee CH_PASS=... \
    .venv/bin/python varvusensor/collector/calibrate.py capture-dark --last 5
"""
from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
from pathlib import Path

import requests

import color
from settings import CH_URL, CH_DB, CH_TABLE, CH_USER, CH_PASS

CALIB_FILE = Path(
    os.environ.get("CALIB_FILE", Path(__file__).with_name("calibration.json"))
)

FIELDS = ("f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8", "clear", "nir")
BANDS = tuple(color.CHANNEL_WAVELENGTHS)  # f1..f8


def ch_query(sql: str) -> str:
    resp = requests.post(
        f"{CH_URL}/",
        params={"database": CH_DB, "query": sql},
        auth=(CH_USER, CH_PASS),
        timeout=15,
    )
    resp.raise_for_status()
    return resp.text


def fetch_rows(last: int) -> list[dict]:
    cols = ", ".join(FIELDS + ("rgb_hex",))
    sql = (
        f"SELECT ts, {cols} FROM {CH_DB}.{CH_TABLE} "
        f"ORDER BY ts DESC LIMIT {int(last)} FORMAT JSONEachRow"
    )
    return [json.loads(line) for line in ch_query(sql).strip().splitlines() if line]


def load_calib() -> dict:
    if not CALIB_FILE.exists():
        return {}
    try:
        return json.loads(CALIB_FILE.read_text())
    except Exception:
        return {}


def save_calib(data: dict) -> None:
    CALIB_FILE.write_text(json.dumps(data, indent=2, ensure_ascii=False))


def capture_reference(kind: str, last: int) -> int:
    rows = fetch_rows(last)
    if not rows:
        print("Andmeid pole — mõõda esmalt (kaetud andur / valge pind).")
        return 1
    ref: dict[str, float] = {}
    for field in FIELDS:
        values = [r[field] for r in rows if r.get(field) is not None]
        ref[field] = round(statistics.mean(values), 3) if values else 0.0
    data = load_calib()
    data[kind] = ref
    data[f"{kind}_rows"] = len(rows)
    save_calib(data)
    label = "Tume" if kind == "dark" else "Valge"
    print(f"{label} referents salvestatud ({len(rows)} reast) -> {CALIB_FILE}")
    print(json.dumps(ref, indent=2))
    if kind == "white" and "dark" not in data:
        print("NB! Tume referents puudub — soovitatav teha enne capture-dark (LED/tumedavool).")
    return 0


def reflectance(
    channels: dict[str, float], dark: dict[str, float], white: dict[str, float]
) -> dict[str, float]:
    out: dict[str, float] = {}
    for band in BANDS:
        value = channels.get(band)
        d = dark.get(band, 0.0) or 0.0
        w = white.get(band, 0.0) or 0.0
        if value is not None and (w - d) > 0:
            out[band] = max(0.0, (value - d) / (w - d))
    return out


def show(last: int) -> int:
    rows = fetch_rows(last)
    data = load_calib()
    dark = data.get("dark") or {}
    white = data.get("white")
    if white is None:
        print("(valge kalibratsioon puudub — näitan ainult toorest; tee capture-white)")
    elif not dark:
        print("(tume referents puudub — kasutan D=0; soovitatav capture-dark)")
    for row in rows:
        channels = {b: row[b] for b in BANDS if row.get(b) is not None}
        raw = color.analyze(channels)
        line = (
            f"{row['ts']}  toores {raw['rgb_hex']} "
            f"(x={raw['x_chrom']:.3f} y={raw['y_chrom']:.3f})"
        )
        if white is not None:
            cal = color.analyze(reflectance(channels, dark, white))
            line += (
                f"  | kalib. {cal['rgb_hex']} "
                f"(x={cal['x_chrom']:.3f} y={cal['y_chrom']:.3f})"
            )
        print(line)
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="AS7341 kalibreerimine")
    sub = parser.add_subparsers(dest="cmd", required=True)
    cd = sub.add_parser("capture-dark", help="salvesta tume referents (andur kaetud)")
    cd.add_argument("--last", type=int, default=5)
    cw = sub.add_parser("capture-white", help="salvesta valge referents")
    cw.add_argument("--last", type=int, default=5)
    sh = sub.add_parser("show", help="näita viimaseid mõõtmisi")
    sh.add_argument("--last", type=int, default=10)
    args = parser.parse_args()
    if args.cmd == "capture-dark":
        return capture_reference("dark", args.last)
    if args.cmd == "capture-white":
        return capture_reference("white", args.last)
    if args.cmd == "show":
        return show(args.last)
    return 1


if __name__ == "__main__":
    sys.exit(main())