05_data_storage.py

"""
05_data_storage.py — BONUS: store IMU/VTS data in queryable formats

Outputs:
  data/sensor_data.db      SQLite  (time-range queries with SQL)
  data/imu_data.parquet    Parquet (columnar, fast range scans)
  data/vts_data.parquet    Parquet

Usage examples:
  python3 05_data_storage.py               # ingest all data
  python3 05_data_storage.py --query       # run example range queries
"""

import os
import sys
import sqlite3
import struct
import argparse
import numpy as np

BASE = os.path.dirname(os.path.abspath(__file__))
IMU_PATH = os.path.join(BASE, 'data', 'recording2.imu')
VTS_PATH = os.path.join(BASE, 'data', 'recording2.vts')
DB_PATH  = os.path.join(BASE, 'data', 'sensor_data.db')
IMU_PQ   = os.path.join(BASE, 'data', 'imu_data.parquet')
VTS_PQ   = os.path.join(BASE, 'data', 'vts_data.parquet')
os.makedirs(os.path.join(BASE, 'data'), exist_ok=True)

sys.path.insert(0, BASE)
from parse_imu_vts_lib import parse_imu, parse_vts, build_sync_index


# ── SQLite ─────────────────────────────────────────────────────────────────────

def ingest_sqlite(imu_records, vts_records, db_path):
    conn = sqlite3.connect(db_path)
    c    = conn.cursor()

    c.executescript("""
        DROP TABLE IF EXISTS imu;
        CREATE TABLE imu (
            id          INTEGER PRIMARY KEY,
            ts_ns       INTEGER NOT NULL,
            ts_us       REAL    NOT NULL,
            ts_s        REAL    NOT NULL,
            accel_x     REAL,
            accel_y     REAL,
            accel_z     REAL,
            gyro_x      REAL,
            gyro_y      REAL,
            gyro_z      REAL,
            mag_x       REAL,
            mag_y       REAL,
            mag_z       REAL,
            temperature REAL
        );
        CREATE INDEX IF NOT EXISTS idx_imu_ts_ns ON imu(ts_ns);
        CREATE INDEX IF NOT EXISTS idx_imu_ts_s  ON imu(ts_s);

        DROP TABLE IF EXISTS vts;
        CREATE TABLE vts (
            id           INTEGER PRIMARY KEY,
            frame_idx    INTEGER NOT NULL,
            imu_idx      INTEGER,
            sync_err_ms  REAL,
            ts_us        INTEGER NOT NULL,
            ts_s         REAL    NOT NULL
        );
        CREATE INDEX IF NOT EXISTS idx_vts_ts_s ON vts(ts_s);
    """)

    imu_rows = [
        (r.timestamp_ns, r.timestamp_us, r.timestamp_s,
         float(r.accel[0]), float(r.accel[1]), float(r.accel[2]),
         float(r.gyro[0]),  float(r.gyro[1]),  float(r.gyro[2]),
         float(r.mag[0]),   float(r.mag[1]),   float(r.mag[2]),
         float(r.temperature))
        for r in imu_records
    ]
    c.executemany(
        "INSERT INTO imu (ts_ns,ts_us,ts_s,accel_x,accel_y,accel_z,"
        "gyro_x,gyro_y,gyro_z,mag_x,mag_y,mag_z,temperature) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)",
        imu_rows
    )

    # Use binary-search-corrected imu_idx (not the raw offset from the binary file)
    sync = build_sync_index(imu_records, vts_records)
    vts_rows = [
        (vr.frame_idx, imu_idx, float(err), vr.timestamp_us, vr.timestamp_s)
        for vr, imu_idx, err in sync
    ]
    c.executemany(
        "INSERT INTO vts (frame_idx,imu_idx,sync_err_ms,ts_us,ts_s) VALUES (?,?,?,?,?)",
        vts_rows
    )

    conn.commit()
    conn.close()
    print(f"[05] SQLite → {db_path}  ({len(imu_rows)} IMU rows, {len(vts_rows)} VTS rows)")


# ── Parquet ────────────────────────────────────────────────────────────────────

def ingest_parquet(imu_records, vts_records, imu_pq, vts_pq):
    try:
        import pandas as pd
        import pyarrow as pa
        import pyarrow.parquet as pq
    except ImportError:
        print("[05] pandas/pyarrow not installed — skipping Parquet. Run: pip install pandas pyarrow")
        return

    # IMU DataFrame
    imu_df = pd.DataFrame({
        'ts_ns'      : [r.timestamp_ns for r in imu_records],
        'ts_s'       : [r.timestamp_s  for r in imu_records],
        'accel_x'    : [float(r.accel[0]) for r in imu_records],
        'accel_y'    : [float(r.accel[1]) for r in imu_records],
        'accel_z'    : [float(r.accel[2]) for r in imu_records],
        'gyro_x'     : [float(r.gyro[0])  for r in imu_records],
        'gyro_y'     : [float(r.gyro[1])  for r in imu_records],
        'gyro_z'     : [float(r.gyro[2])  for r in imu_records],
        'mag_x'      : [float(r.mag[0])   for r in imu_records],
        'mag_y'      : [float(r.mag[1])   for r in imu_records],
        'mag_z'      : [float(r.mag[2])   for r in imu_records],
        'temperature': [float(r.temperature) for r in imu_records],
    })
    # Sort and write with snappy compression + row group size for efficient range scans
    imu_df.sort_values('ts_s', inplace=True)
    table = pa.Table.from_pandas(imu_df, preserve_index=False)
    pq.write_table(table, imu_pq, compression='snappy', row_group_size=5000)
    print(f"[05] Parquet → {imu_pq}  ({len(imu_df)} rows)")

    # VTS DataFrame — use binary-search-corrected imu_idx, not raw binary offset
    sync = build_sync_index(imu_records, vts_records)
    correct_imu_idx = [idx for _, idx, _ in sync]
    sync_errors     = [err for _, _, err in sync]
    vts_df = pd.DataFrame({
        'frame_idx'  : [r.frame_idx     for r in vts_records],
        'imu_idx'    : correct_imu_idx,   # binary-search corrected
        'sync_err_ms': sync_errors,        # frame↔IMU timestamp delta in ms
        'ts_us'      : [r.timestamp_us for r in vts_records],
        'ts_s'       : [r.timestamp_s  for r in vts_records],
    })
    vts_df.sort_values('ts_s', inplace=True)
    table_vts = pa.Table.from_pandas(vts_df, preserve_index=False)
    pq.write_table(table_vts, vts_pq, compression='snappy')
    print(f"[05] Parquet → {vts_pq}  ({len(vts_df)} rows)")


# ── Example queries ────────────────────────────────────────────────────────────

def run_example_queries(db_path):
    import sqlite3
    conn = sqlite3.connect(db_path)
    c    = conn.cursor()

    t_start = 40.0
    t_end   = 42.0
    print(f"\n── SQL: IMU data in t=[{t_start}, {t_end}]s ─────────────────────")
    c.execute("""
        SELECT ts_s, accel_x, accel_y, accel_z, gyro_x, gyro_y, gyro_z, temperature
        FROM imu
        WHERE ts_s BETWEEN ? AND ?
        LIMIT 10
    """, (t_start, t_end))
    rows = c.fetchall()
    print(f"  Got {len(rows)} sample rows (showing first 5):")
    for row in rows[:5]:
        print(f"    t={row[0]:.4f}s  accel=({row[1]:.3f},{row[2]:.3f},{row[3]:.3f})"
              f"  gyro=({row[4]:.4f},{row[5]:.4f},{row[6]:.4f})  temp={row[7]:.2f}°C")

    print(f"\n── SQL: avg accel magnitude in t=[{t_start}, {t_end}]s ──────────")
    c.execute("""
        SELECT
            AVG(SQRT(accel_x*accel_x + accel_y*accel_y + accel_z*accel_z)) AS avg_mag,
            AVG(temperature) AS avg_temp,
            COUNT(*) AS n_samples
        FROM imu
        WHERE ts_s BETWEEN ? AND ?
    """, (t_start, t_end))
    row = c.fetchone()
    print(f"  avg |a|={row[0]:.4f} m/s²  avg_temp={row[1]:.2f}°C  n={row[2]}")

    print(f"\n── SQL: video frames in t=[{t_start}, {t_end}]s ─────────────────")
    c.execute("""
        SELECT frame_idx, ts_s FROM vts
        WHERE ts_s BETWEEN ? AND ?
        ORDER BY frame_idx
    """, (t_start, t_end))
    vts_rows = c.fetchall()
    print(f"  {len(vts_rows)} frames: idx {vts_rows[0][0]}..{vts_rows[-1][0]}" if vts_rows else "  no frames")

    conn.close()


# ── Parquet range query example ────────────────────────────────────────────────

def run_parquet_query(imu_pq, t_start=40.0, t_end=42.0):
    try:
        import pandas as pd
        import pyarrow.parquet as pq
        import pyarrow.compute as pc

        print(f"\n── Parquet: IMU in t=[{t_start}, {t_end}]s ──────────────────")
        table = pq.read_table(imu_pq, filters=[
            ('ts_s', '>=', t_start),
            ('ts_s', '<=', t_end),
        ])
        df = table.to_pandas()
        print(f"  {len(df)} records loaded from Parquet")
        print(f"  accel_z mean: {df['accel_z'].mean():.4f} m/s²")
        print(f"  temperature : {df['temperature'].mean():.2f}°C")
    except ImportError:
        print("[05] pandas/pyarrow not available for Parquet query.")


# ── Main ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Ingest sensor data to SQLite/Parquet")
    parser.add_argument('--query', action='store_true', help="Run example range queries")
    args = parser.parse_args()

    print("[05] Parsing binary files …")
    imu_records = parse_imu(IMU_PATH)
    vts_records = parse_vts(VTS_PATH)

    print("[05] Ingesting to SQLite …")
    ingest_sqlite(imu_records, vts_records, DB_PATH)

    print("[05] Ingesting to Parquet …")
    ingest_parquet(imu_records, vts_records, IMU_PQ, VTS_PQ)

    if args.query or True:   # always run queries for demo
        run_example_queries(DB_PATH)
        run_parquet_query(IMU_PQ)

    print("\n[05] Done.")


if __name__ == '__main__':
    main()