# gui/components/data/decode_raw_data.py
# =============================================================================
# Purpose
# -------
# Decode raw rows from **one session** in 'telemetry' into 'decoded_telemetry'.
#
# Scope
# -----
# - This decoder processes **only a single session_id**.
# - If session_id is not provided, it auto-selects the **most recent** session
#   present in 'telemetry' (by MAX(t_ns)).
#
# What it does
# ------------
# - Reads rows from telemetry for the chosen session_id.
# - Copies session metadata **exactly** from telemetry:
#     session_id, logger_session, logger_session_number.
# - Decodes UART 14B frames (EF FE ... EE), Big Endian pairs:
#     [2..3]=u16 MotorCurrent
#     [4..5]=s16 EncoderValue
#     [6..7]=s16 RelativeEncoderValue
#     [8..9]=s16 V24PecDiff
#     [10..11]=s16 Reserved0
#     [12]=u8  PWM
# - Computes time_ms per (session_id, logger_session_number) relative to first t_ns.
# - Parses I2C 2B angle (AS5048A) into:
#     i2c_raw14, i2c_zero_raw14 (first seen per (session_id, logger_session_number)),
#     i2c_delta_raw14 (signed wrap),
#     i2c_angle_deg (-(delta*360/16384)),
#     i2c_zero_angle_deg (zero*360/16384)
# - Upserts decoded rows into 'decoded_telemetry' (no i2c_raw stored).
#
# Entry points
# ------------
# - decode_into_decoded_telemetry(session_id: str|None = None, batch_size: int = 1000) -> dict
# - run(session_id: str|None = None, **_) -> dict          (alias; ignores other filters)
# - main([session_id]) -> int                               (simple CLI)
# =============================================================================

from __future__ import annotations
from typing import Optional, Tuple, Dict, Any, List

# DB helpers (centralized)
from components.data.db import (
    get_connection,
    ensure_decoded_schema,
    upsert_decoded_rows,
)

# -----------------------------
# UART framing & decode helpers
# -----------------------------
FRAME_LEN = 14
START0 = 0xEF
START1 = 0xFE
END = 0xEE


def _parse_hex_bytes(s: Optional[str]) -> Optional[bytes]:
    """Accepts 'EF FE ... EE' or 'EFFEEE...' (spaces optional). Returns None if invalid/missing."""
    if not s:
        return None
    s2 = s.strip().replace(" ", "")
    if len(s2) % 2 != 0:
        return None
    try:
        return bytes.fromhex(s2)
    except ValueError:
        return None


def _be_u16(b0: int, b1: int) -> int:
    return (b0 << 8) | b1


def _be_s16(b0: int, b1: int) -> int:
    u = (b0 << 8) | b1
    return u - 0x10000 if (u & 0x8000) else u


def _decode_uart_frame(
    raw: Optional[bytes],
) -> Optional[Tuple[int, int, int, int, int, int]]:
    if raw is None or len(raw) != FRAME_LEN:
        return None
    if not (raw[0] == START0 and raw[1] == START1 and raw[13] == END):
        return None

    # motor_current must be signed s16 (two's complement, big endian)
    mc = _be_s16(raw[2], raw[3])
    enc = _be_s16(raw[4], raw[5])
    rel = _be_s16(raw[6], raw[7])
    v24 = _be_s16(raw[8], raw[9])
    res0 = _be_s16(raw[10], raw[11])
    pwm = int(raw[12])  # u8 (0..100)

    # Optional sanity guard to drop obvious misframes (tune bounds to your hw)
    if not (-40000 <= mc <= 40000 and 0 <= pwm <= 100):
        return None

    return (mc, enc, rel, v24, res0, pwm)


# -----------------------------
# I2C decode helpers (AS5048A)
# -----------------------------
def _i2c_raw14(i2c_raw: Optional[str]) -> Optional[int]:
    """Decode 2-byte I2C payload into 14-bit angle (0..16383)."""
    b = _parse_hex_bytes(i2c_raw)
    if b is None or len(b) != 2:
        return None
    msb, lsb = b[0], b[1]
    return ((msb << 6) | (lsb & 0x3F)) & 0x3FFF


def _wrap14_delta(cur: int, zero: int) -> int:
    """Signed wrap-around delta (cur - zero) mapped to [-8192..+8191] in 14-bit space."""
    return ((cur - zero + 8192) % 16384) - 8192


# -----------------------------
# Session selection helpers
# -----------------------------
def _resolve_session_id() -> Optional[str]:
    """
    Pick the most recent session_id present in telemetry (by MAX(t_ns)).
    Returns None if telemetry is empty.
    """
    with get_connection() as con:
        cur = con.cursor()
        cur.execute(
            """
            SELECT session_id
            FROM telemetry
            WHERE t_ns IS NOT NULL
            GROUP BY session_id
            ORDER BY MAX(t_ns) DESC
            LIMIT 1
            """
        )
        row = cur.fetchone()
        return row[0] if row else None


# -----------------------------
# Internal: load rows for one session
# -----------------------------
_SELECT_FOR_SESSION = """
SELECT
  id, session_id, logger_session, logger_session_number,
  t_ns, uart_raw, i2c_raw
FROM telemetry
WHERE session_id = ?
ORDER BY logger_session_number, t_ns
"""


# -----------------------------
# Public: main decode entry
# -----------------------------
def decode_into_decoded_telemetry(
    session_id: Optional[str] = None,
    *,
    batch_size: int = 1000,
) -> Dict[str, int]:
    ensure_decoded_schema()

    sess = session_id or _resolve_session_id()
    if not sess:
        return {"processed": 0, "decoded": 0, "skipped_invalid_uart": 0}

    # Load rows for this session
    with get_connection() as con:
        cur = con.cursor()
        cur.execute(_SELECT_FOR_SESSION, (sess,))
        rows = cur.fetchall()

    processed = decoded = skipped = 0

    first_ns_map: Dict[Tuple[str, int], int] = {}
    i2c_zero_map: Dict[Tuple[str, int], int] = {}

    try:
        import config

        _CONFIG_ZERO14 = (
            ((int(config.SESSION_ZERO[0]) & 0xFF) << 8)
            | (int(config.SESSION_ZERO[1]) & 0xFF)
        ) & 0x3FFF
    except Exception:
        _CONFIG_ZERO14 = None

    # 1) PASS: decode & COLLECT per run
    from collections import defaultdict

    grp_rows: Dict[Tuple[str, int], List[Dict[str, Any]]] = defaultdict(list)

    for id_, sess_id, log_tag, log_no, t_ns, uart_raw, i2c_raw in rows:
        processed += 1
        grp = (sess_id, int(log_no))

        if _CONFIG_ZERO14 is not None:
            i2c_zero_map.setdefault(grp, _CONFIG_ZERO14)

        # baseline for time_ms
        if grp not in first_ns_map and t_ns is not None:
            first_ns_map[grp] = t_ns
        t0 = first_ns_map.get(grp)
        time_ms = (
            int((t_ns - t0) / 1_000_000)
            if (t0 is not None and t_ns is not None)
            else None
        )

        # UART decode
        frame = _decode_uart_frame(_parse_hex_bytes(uart_raw))
        if frame is None:
            skipped += 1
            continue
        mc, enc, rel, v24, res0, pwm = frame

        # I2C decode/angles
        raw14 = _i2c_raw14(i2c_raw)
        zero14 = i2c_zero_map.get(grp)
        if raw14 is not None and zero14 is None:
            zero14 = raw14
            i2c_zero_map[grp] = zero14

        if raw14 is not None and zero14 is not None:
            delta14 = _wrap14_delta(raw14, zero14)
            angle_deg = -(delta14 * 360.0 / 16384.0)
            zero_angle_deg = zero14 * 360.0 / 16384.0
        else:
            delta14 = None
            angle_deg = None
            zero_angle_deg = (zero14 * 360.0 / 16384.0) if zero14 is not None else None

        grp_rows[grp].append(
            {
                "id": id_,
                "session_id": sess_id,
                "logger_session": log_tag,
                "logger_session_number": int(log_no),
                "t_ns": t_ns,
                "time_ms": time_ms,
                "motor_current": mc,
                "encoder_value": enc,
                "relative_encoder_value": rel,
                "v24_pec_diff": v24,
                "pwm": pwm,
                "i2c_raw14": raw14,
                "i2c_zero_raw14": zero14,
                "i2c_delta_raw14": delta14,
                "i2c_angle_deg": angle_deg,
                "i2c_zero_angle_deg": zero_angle_deg,
                "angular_velocity": None,
            }
        )

    # 2) PASS: compute angular_velocity per group using Kalman RTS on angle vs time
    for grp, items in grp_rows.items():
        # Build time in seconds with full precision; keep your existing time_ms for storage
        time_s_list = []
        angle_deg_list = []
        for row in items:
            t_ns = row["t_ns"]
            # use high-precision seconds internally
            time_s_list.append(
                (t_ns - items[0]["t_ns"]) * 1e-9
                if (t_ns is not None and items[0]["t_ns"] is not None)
                else None
            )
            angle_deg_list.append(
                row["i2c_angle_deg"]
            )  # may be None; RTS handles segments

        omega_deg_s = _kalman_velocity_over_series(
            time_s_list, angle_deg_list, gap_threshold_s=0.1
        )

        # write back
        for row, w in zip(items, omega_deg_s):
            row["angular_velocity"] = w

    # 3) UPSERT in batches (same API)
    out_batch: List[Dict[str, Any]] = []

    def _flush_batch():
        nonlocal out_batch, decoded
        if not out_batch:
            return
        decoded += upsert_decoded_rows(out_batch)
        out_batch = []

    for items in grp_rows.values():
        for row in items:
            out_batch.append(row)
            if len(out_batch) >= batch_size:
                _flush_batch()
    _flush_batch()

    return {"processed": processed, "decoded": decoded, "skipped_invalid_uart": skipped}


# -----------------------------
# Kalman RTS (constant-velocity) helpers
# -----------------------------
def _robust_var(values):
    """Robust variance via MAD; returns a small floor if data is tiny."""
    xs = [v for v in values if v is not None]
    if len(xs) < 5:
        return 1e-6
    xs_sorted = sorted(xs)
    m = xs_sorted[len(xs) // 2]
    abs_dev = [abs(v - m) for v in xs_sorted]
    mad = sorted(abs_dev)[len(abs_dev) // 2] or 1e-9
    # 1.4826 * MAD ~ sigma (Gaussian)
    sigma = 1.4826 * mad
    v = sigma * sigma
    return v if v > 1e-6 else 1e-6


def _kalman_rts_velocity_segment(t_s, y_deg, r_var=None, q_pow=None):
    """
    One contiguous segment (strictly increasing times, no big gaps).
    State x=[theta, omega], units: theta=deg, omega=deg/s.
    F_k = [[1, dt],[0,1]]
    Q_k = q * [[dt^3/3, dt^2/2],[dt^2/2, dt]]
    H   = [1, 0], R = r.
    Returns (theta_s, omega_s) lists (len = len(y_deg))
    """
    n = len(y_deg)
    if n == 0:
        return [], []
    # Estimate R (measurement variance) if not provided
    if r_var is None:
        r_var = _robust_var(y_deg)
    # Estimate process noise power q if not provided (based on rough accel scale)
    if q_pow is None:
        # crude accel estimate from finite differences
        dtheta = [(y_deg[i + 1] - y_deg[i]) for i in range(n - 1)]
        dt = [(t_s[i + 1] - t_s[i]) for i in range(n - 1)]
        dv = []
        for i in range(len(dtheta) - 1):
            dt1 = dt[i] if dt[i] > 1e-6 else 1e-6
            dt2 = dt[i + 1] if dt[i + 1] > 1e-6 else 1e-6
            v1 = dtheta[i] / dt1
            v2 = dtheta[i + 1] / dt2
            dv.append(
                (v2 - v1) / max((dt1 + dt2) / 2.0, 1e-6)
            )  # deg/s^2 change per second
        # variance-like scale of acceleration
        acc_var = _robust_var(dv) if dv else 1.0
        med_dt = sorted(dt)[len(dt) // 2] if dt else 0.01
        q_pow = max(acc_var * med_dt * 1e-3, 1e-6)  # conservative floor

    # Allocate
    x_f = [[0.0, 0.0] for _ in range(n)]
    P_f = [[[0.0, 0.0], [0.0, 0.0]] for _ in range(n)]
    x_p = [[0.0, 0.0] for _ in range(n)]
    P_p = [[[0.0, 0.0], [0.0, 0.0]] for _ in range(n)]

    # Init
    x = [y_deg[0], 0.0]
    P = [[1.0, 0.0], [0.0, 1000.0]]

    H = [1.0, 0.0]
    R = r_var

    def mat_add(A, B):
        return [
            [A[0][0] + B[0][0], A[0][1] + B[0][1]],
            [A[1][0] + B[1][0], A[1][1] + B[1][1]],
        ]

    def mat_mul(A, B):
        return [
            [
                A[0][0] * B[0][0] + A[0][1] * B[1][0],
                A[0][0] * B[0][1] + A[0][1] * B[1][1],
            ],
            [
                A[1][0] * B[0][0] + A[1][1] * B[1][0],
                A[1][0] * B[0][1] + A[1][1] * B[1][1],
            ],
        ]

    def mat_transpose(A):
        return [[A[0][0], A[1][0]], [A[0][1], A[1][1]]]

    def mat_vec(A, v):
        return [A[0][0] * v[0] + A[0][1] * v[1], A[1][0] * v[0] + A[1][1] * v[1]]

    def mat_inv_2x2(A):
        det = A[0][0] * A[1][1] - A[0][1] * A[1][0]
        if abs(det) < 1e-12:
            det = 1e-12
        inv = [[A[1][1] / det, -A[0][1] / det], [-A[1][0] / det, A[0][0] / det]]
        return inv

    # FILTER
    for k in range(n):
        if k == 0:
            x_p[k] = [x[0], x[1]]
            P_p[k] = [[P[0][0], P[0][1]], [P[1][0], P[1][1]]]
        else:
            dt = t_s[k] - t_s[k - 1]
            F = [[1.0, dt], [0.0, 1.0]]
            q11 = (dt**3) / 3.0
            q12 = (dt**2) / 2.0
            q22 = dt
            Q = [[q_pow * q11, q_pow * q12], [q_pow * q12, q_pow * q22]]
            # Predict
            x = [x[0] + dt * x[1], x[1]]
            FP = mat_mul(F, P)
            FT = mat_transpose(F)
            P = mat_add(mat_mul(FP, FT), Q)
            x_p[k] = [x[0], x[1]]
            P_p[k] = [[P[0][0], P[0][1]], [P[1][0], P[1][1]]]

        # Update with measurement y_deg[k]
        z = y_deg[k]
        # S = H*P*H^T + R  (scalar)
        S = P[0][0] + R
        K0 = P[0][0] / S  # first row, col 0 of K
        K1 = P[1][0] / S  # second row, col 0 of K
        innov = z - x[0]
        x = [x[0] + K0 * innov, x[1] + K1 * innov]
        # P = (I - K*H)P
        # (I - K H) = [[1-K0*1, -K0*0],[-K1*1, 1- K1*0]] = [[1-K0, 0],[-K1, 1]]
        IKH = [[1.0 - K0, 0.0], [-K1, 1.0]]
        P = mat_mul(IKH, P)

        x_f[k] = [x[0], x[1]]
        P_f[k] = [[P[0][0], P[0][1]], [P[1][0], P[1][1]]]

    # RTS SMOOTHER
    x_s = [[x_f[i][0], x_f[i][1]] for i in range(n)]
    P_s = [
        [[P_f[i][0][0], P_f[i][0][1]], [P_f[i][1][0], P_f[i][1][1]]] for i in range(n)
    ]

    for k in range(n - 2, -1, -1):
        dt = t_s[k + 1] - t_s[k]
        F = [[1.0, dt], [0.0, 1.0]]
        FT = mat_transpose(F)
        Pp_inv = mat_inv_2x2(P_p[k + 1])
        # Ck = P_f[k] * F^T * P_p[k+1]^{-1}
        C_left = mat_mul(P_f[k], FT)
        Ck = mat_mul(C_left, Pp_inv)
        # x_s[k] = x_f[k] + Ck * (x_s[k+1] - x_p[k+1])
        dx = [x_s[k + 1][0] - x_p[k + 1][0], x_s[k + 1][1] - x_p[k + 1][1]]
        corr = mat_vec(Ck, dx)
        x_s[k] = [x_f[k][0] + corr[0], x_f[k][1] + corr[1]]
        # P_s[k] not needed for velocity output, so we skip computing it for speed

    theta_s = [x_s[i][0] for i in range(n)]
    omega_s = [x_s[i][1] for i in range(n)]
    return theta_s, omega_s


def _kalman_velocity_over_series(
    time_s_list, angle_deg_list, *, gap_threshold_s=0.1
) -> list:
    """
    Splits the series at big gaps or invalid points and runs the RTS smoother per segment.
    Returns a list of angular velocities (deg/s) aligned with inputs (None where undefined).
    """
    n = len(angle_deg_list)
    out = [None] * n
    # Build valid indices where both time and angle exist
    valid = [
        (time_s_list[i] is not None) and (angle_deg_list[i] is not None)
        for i in range(n)
    ]
    i = 0
    while i < n:
        # skip invalids
        while i < n and not valid[i]:
            i += 1
        if i >= n:
            break
        j = i + 1
        # extend segment until gap or invalid
        while j < n and valid[j]:
            if (time_s_list[j] - time_s_list[j - 1]) > gap_threshold_s:
                break
            j += 1
        # segment [i, j)
        t_seg = [time_s_list[k] for k in range(i, j)]
        y_seg = [angle_deg_list[k] for k in range(i, j)]
        theta_s, omega_s = _kalman_rts_velocity_segment(t_seg, y_seg)
        for k, w in enumerate(omega_s):
            out[i + k] = w
        i = j
    return out


# Friendly alias for UART_logic.stop_logger()
def run(session_id: Optional[str] = None, **_ignored) -> Dict[str, int]:
    """
    Alias to decode_into_decoded_telemetry with **only** session_id honored.
    Other kwargs are ignored intentionally.
    """
    return decode_into_decoded_telemetry(session_id=session_id)


def main() -> int:
    """CLI: python -m gui.components.data.decode_raw_data [session_id]"""
    import sys

    sess = sys.argv[1] if len(sys.argv) > 1 else None
    stats = decode_into_decoded_telemetry(session_id=sess)
    print(stats)
    return 0


if __name__ == "__main__":
    main()