ohmu/dev/varvusensor/collector/calibrate.py

"""Kalibreerimise abivahend AS7341 värviandurile.

Loeb ClickHouse'ist (vaikimisi uus67.varvus) ja toetab:
  * capture-white: salvesta valge referents (viimased N mõõtmist) -> calibration.json
  * show:          näita viimaseid mõõtmisi (toores hex vs kalibreeritud hex)

Kalibreeritud värvus = iga kanal jagatud valge referentsiga (reflectance),
mis eemaldab valgustuse värvitooni ja anduri kanalite erinevuse.

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-white --last 5
  CH_URL=http://100.87.1.24:8125 CH_USER=zigbee CH_PASS=... \
    .venv/bin/python varvusensor/collector/calibrate.py show --last 10
"""
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 capture_white(last: int) -> int:
    rows = fetch_rows(last)
    if not rows:
        print("Andmeid pole — mõõda esmalt valget pinda.")
        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
    payload = {"reference_rows": len(rows), "white": ref}
    CALIB_FILE.write_text(json.dumps(payload, indent=2, ensure_ascii=False))
    print(f"Valge referents salvestatud ({len(rows)} reast) -> {CALIB_FILE}")
    print(json.dumps(ref, indent=2))
    return 0


def load_white() -> dict[str, float] | None:
    if not CALIB_FILE.exists():
        return None
    return json.loads(CALIB_FILE.read_text()).get("white")


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


def show(last: int) -> int:
    rows = fetch_rows(last)
    white = load_white()
    if white is None:
        print("(kalibratsioon puudub — näitan ainult toorest; tee capture-white)")
    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, 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)
    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-white":
        return capture_white(args.last)
    if args.cmd == "show":
        return show(args.last)
    return 1


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