#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
record_episode.py

PC-side episode recorder.

Records:
  - ESP32 IMU data: acc / gyro / angle / temperature
  - ESP32 DIN state: GPIO25 / GPIO26 / GPIO27 / GPIO32
  - ESP32 DOUT logical state: GPIO18 / GPIO19 relay command state
  - inferred scissor state: open / close / stop
  - RealSense D405 RGB image
  - RealSense D405 depth image
  - Filtered IMU integrated velocity

Default DIN mapping from ESP32 main.py:
  DIN_PINS = (25, 26, 27, 32)

Therefore:
  din[0] = GPIO25
  din[1] = GPIO26
  din[2] = GPIO27   # already used as physical scissor switch
  din[3] = GPIO32

Default recording trigger:
  din[0] / GPIO25

Usage:

  Test ESP32 switch + CSV only:
    python3 record_episode.py --port /dev/ttyUSB0 --no-camera

  Record with D405 RGB + depth:
    python3 record_episode.py --port /dev/ttyUSB0

  Use GPIO26 as record trigger:
    python3 record_episode.py --port /dev/ttyUSB0 --record-din 1
"""

import argparse
import csv
import time
import math
from pathlib import Path
from datetime import datetime

import cv2
import numpy as np

try:
    import pyrealsense2 as rs
except ImportError:
    rs = None

from esp_bridge import ESP32Bridge


# ============================================================
# Switch edge detector
# ============================================================

class SwitchChangeDetector:
    """
    Detect any valid state change from a toggle switch.

    For a latching/toggle switch:
      1 -> 0 triggers once
      0 -> 1 triggers once

    This is different from a push button, where only "pressed" edge is used.
    """

    def __init__(self, debounce_s=0.25):
        self.debounce_s = debounce_s
        self.last_value = None
        self.last_edge_time = 0.0

    def update(self, raw_value):
        now = time.monotonic()

        # First sample only initializes state.
        # Do not trigger recording immediately when program starts.
        if self.last_value is None:
            self.last_value = raw_value
            return False

        if raw_value == self.last_value:
            return False

        if now - self.last_edge_time < self.debounce_s:
            return False

        old_value = self.last_value
        self.last_value = raw_value
        self.last_edge_time = now

        print(f"[SWITCH CHANGE] {old_value} -> {raw_value}")

        return True

# ============================================================
# Switch level controller
# ============================================================
class SwitchLevelController:
    """
    Level-control recording switch.

    raw_value == active_value -> should record
    raw_value != active_value -> should stop

    For your current requirement:
      DIN0 = 1 -> REC
      DIN0 = 0 -> IDLE
    """

    def __init__(self, active_value=1, debounce_s=0.25):
        self.active_value = active_value
        self.debounce_s = debounce_s

        self.last_raw_value = None
        self.stable_value = None
        self.last_change_time = 0.0

    def update(self, raw_value):
        now = time.monotonic()

        # First sample initializes stable state.
        # If program starts while switch is already 1, it will start recording.
        if self.last_raw_value is None:
            self.last_raw_value = raw_value
            self.stable_value = raw_value
            return raw_value == self.active_value

        # Raw value changed. Start debounce timer.
        if raw_value != self.last_raw_value:
            self.last_raw_value = raw_value
            self.last_change_time = now
            return None

        # Raw value has remained the same. Accept it after debounce time.
        if raw_value != self.stable_value:
            if now - self.last_change_time >= self.debounce_s:
                old_value = self.stable_value
                self.stable_value = raw_value
                print(f"[SWITCH LEVEL] {old_value} -> {self.stable_value}")
                return self.stable_value == self.active_value

        return None

# ============================================================
# RealSense D405 camera
# ============================================================

class D405Camera:
    def __init__(
        self,
        width=640,
        height=480,
        fps=30,
        enable_depth=True,
        warmup_frames=15,
    ):
        if rs is None:
            raise RuntimeError(
                "pyrealsense2 is not installed. "
                "Install it or run with --no-camera."
            )

        self.width = width
        self.height = height
        self.fps = fps
        self.enable_depth = enable_depth
        self.warmup_frames = warmup_frames

        self.pipeline = rs.pipeline()
        self.config = rs.config()
        self.align = None

    def open(self):
        self.config.enable_stream(
            rs.stream.color,
            self.width,
            self.height,
            rs.format.bgr8,
            self.fps,
        )

        if self.enable_depth:
            self.config.enable_stream(
                rs.stream.depth,
                self.width,
                self.height,
                rs.format.z16,
                self.fps,
            )
            self.align = rs.align(rs.stream.color)

        self.pipeline.start(self.config)

        # Warm up auto exposure.
        for _ in range(self.warmup_frames):
            self.pipeline.wait_for_frames()

        return self

    def read(self):
        frames = self.pipeline.wait_for_frames()

        if self.enable_depth and self.align is not None:
            frames = self.align.process(frames)

        color_frame = frames.get_color_frame()
        if not color_frame:
            return None, None

        color = np.asanyarray(color_frame.get_data())

        depth = None
        if self.enable_depth:
            depth_frame = frames.get_depth_frame()
            if depth_frame:
                depth = np.asanyarray(depth_frame.get_data())

        return color, depth

    def close(self):
        try:
            self.pipeline.stop()
        except Exception:
            pass

    def __enter__(self):
        return self.open()

    def __exit__(self, exc_type, exc, tb):
        self.close()


# ============================================================
# Simple IMU velocity estimator
# ============================================================

class ProviderIMUVelocityEstimator:
    """
    Provider-style roll/pitch gravity compensation + velocity integration.

    This follows verify.py:
      gravity_body_g = [
          -sin(pitch),
          cos(pitch) * sin(roll),
          cos(pitch) * cos(roll),
      ] * gravity_sign

    Then:
      linear_acc_body = raw_acc_body - gravity_body

    Important:
      This is still an estimated IMU velocity, not ground truth.
      It removes the static gravity/support acceleration approximately,
      but velocity integration will still drift.
    """

    def __init__(
        self,
        gravity_sign=1.0,
        bias_seconds=1.0,
        acc_deadband=0.10,
        vel_decay=0.995,
    ):
        self.gravity_sign = gravity_sign
        self.bias_seconds = bias_seconds
        self.acc_deadband = acc_deadband
        self.vel_decay = vel_decay

        self.start_t_ms = None
        self.last_t_ms = None

        self.bias_ready = False
        self.bias_samples = []
        self.acc_bias = np.zeros(3, dtype=float)

        self.v = np.zeros(3, dtype=float)

    def reset(self):
        self.start_t_ms = None
        self.last_t_ms = None

        self.bias_ready = False
        self.bias_samples = []
        self.acc_bias[:] = 0.0

        self.v[:] = 0.0

    def gravity_body_g(self, roll_deg, pitch_deg):
        roll = math.radians(roll_deg)
        pitch = math.radians(pitch_deg)

        g_x = -math.sin(pitch)
        g_y = math.cos(pitch) * math.sin(roll)
        g_z = math.cos(pitch) * math.cos(roll)

        return np.array([g_x, g_y, g_z], dtype=float) * self.gravity_sign

    def update(self, packet):
        t_ms = packet["t_ms"]

        if self.start_t_ms is None:
            self.start_t_ms = t_ms

        acc_body_g = np.array(packet["acc_g"], dtype=float)

        roll_deg = packet["angle_deg"][0]
        pitch_deg = packet["angle_deg"][1]

        gravity_g = self.gravity_body_g(
            roll_deg=roll_deg,
            pitch_deg=pitch_deg,
        )

        linear_acc_g_raw = acc_body_g - gravity_g
        linear_acc_ms2_raw = linear_acc_g_raw * 9.81

        elapsed_s = ((t_ms - self.start_t_ms) & 0xFFFFFFFF) / 1000.0

        # First N seconds are used as residual bias estimation.
        # During this period velocity stays zero.
        if not self.bias_ready:
            self.bias_samples.append(linear_acc_ms2_raw)

            if elapsed_s >= self.bias_seconds:
                if self.bias_samples:
                    self.acc_bias = np.mean(
                        np.array(self.bias_samples, dtype=float),
                        axis=0,
                    )

                self.bias_ready = True
                print(f"[IMU] Provider acc bias estimated: {self.acc_bias}")

            self.last_t_ms = t_ms

            linear_acc_ms2 = linear_acc_ms2_raw - self.acc_bias

            return {
                "gravity_g": gravity_g.copy(),
                "linear_acc_ms2": linear_acc_ms2.copy(),
                "velocity_ms": self.v.copy(),
                "acc_bias_ms2": self.acc_bias.copy(),
                "bias_ready": self.bias_ready,
            }

        delta_ms = (t_ms - self.last_t_ms) & 0xFFFFFFFF
        self.last_t_ms = t_ms

        linear_acc_ms2 = linear_acc_ms2_raw - self.acc_bias

        if np.linalg.norm(linear_acc_ms2) < self.acc_deadband:
            linear_acc_ms2[:] = 0.0

        if 0 < delta_ms <= 200:
            dt = delta_ms / 1000.0
            self.v = self.v * self.vel_decay + linear_acc_ms2 * dt

        return {
            "gravity_g": gravity_g.copy(),
            "linear_acc_ms2": linear_acc_ms2.copy(),
            "velocity_ms": self.v.copy(),
            "acc_bias_ms2": self.acc_bias.copy(),
            "bias_ready": self.bias_ready,
        }


# ============================================================
# Scissor state
# ============================================================

def infer_scissor_state(packet):
    """
    DOUT logical mapping:

      dout[0] = GPIO18 = relay IN1
      dout[1] = GPIO19 = relay IN2

    Hardware behavior:
      IN1 active -> CLOSE
      IN2 active -> OPEN

    Important:
      packet["dout"] comes from ESP32 dout_shadow.
      It is a logical active state, not the raw GPIO voltage.
    """

    dout0 = packet["dout"][0]
    dout1 = packet["dout"][1]

    if dout0 == 1 and dout1 == 0:
        return "close"

    if dout0 == 0 and dout1 == 1:
        return "open"

    if dout0 == 0 and dout1 == 0:
        return "stop"

    return "invalid_both_active"


# ============================================================
# Episode writer
# ============================================================

class EpisodeWriter:
    def __init__(
        self,
        root_dir="dataset",
        save_rgb=True,
        save_depth=True,
        jpeg_quality=95,
        flush_every=10,
    ):
        self.root_dir = Path(root_dir)
        self.save_rgb = save_rgb
        self.save_depth = save_depth
        self.jpeg_quality = jpeg_quality
        self.flush_every = flush_every

        self.root_dir.mkdir(parents=True, exist_ok=True)

        self.recording = False

        self.episode_dir = None
        self.rgb_dir = None
        self.depth_dir = None

        self.csv_file = None
        self.writer = None

        self.frame_id = 0

    def start(self):
        if self.recording:
            return

        episode_name = "episode_" + datetime.now().strftime("%Y%m%d_%H%M%S")
        self.episode_dir = self.root_dir / episode_name

        self.rgb_dir = self.episode_dir / "rgb"
        self.depth_dir = self.episode_dir / "depth"

        self.episode_dir.mkdir(parents=True, exist_ok=True)

        if self.save_rgb:
            self.rgb_dir.mkdir(parents=True, exist_ok=True)

        if self.save_depth:
            self.depth_dir.mkdir(parents=True, exist_ok=True)

        csv_path = self.episode_dir / "data.csv"
        self.csv_file = open(csv_path, "w", newline="", encoding="utf-8")

        fieldnames = [
            "frame_id",

            "t_wall",
            "t_monotonic",
            "t_esp_ms",

            "din0_gpio25_record_switch",
            "din1_gpio26",
            "din2_gpio27_scissor_switch",
            "din3_gpio32",

            "dout0_gpio18_in1",
            "dout1_gpio19_in2",
            "scissor_state",

            "acc_x_g",
            "acc_y_g",
            "acc_z_g",

            "gyro_x_dps",
            "gyro_y_dps",
            "gyro_z_dps",

            "angle_x_deg",
            "angle_y_deg",
            "angle_z_deg",

            "temp_c",

            "vel_x_ms",
            "vel_y_ms",
            "vel_z_ms",

            "lin_acc_x_ms2",
            "lin_acc_y_ms2",
            "lin_acc_z_ms2",

            "acc_bias_x_ms2",
            "acc_bias_y_ms2",
            "acc_bias_z_ms2",

            "bias_ready",

            "rgb_path",
            "depth_path",
        ]

        self.writer = csv.DictWriter(self.csv_file, fieldnames=fieldnames)
        self.writer.writeheader()

        self.frame_id = 0
        self.recording = True

        print(f"[REC START] {self.episode_dir}")

    def stop(self):
        if not self.recording:
            return

        self.recording = False

        if self.csv_file is not None:
            self.csv_file.flush()
            self.csv_file.close()

        print(f"[REC STOP] {self.episode_dir}")

        self.csv_file = None
        self.writer = None

        self.episode_dir = None
        self.rgb_dir = None
        self.depth_dir = None

        self.frame_id = 0

    def toggle(self):
        if self.recording:
            self.stop()
        else:
            self.start()

    def write(self, packet, imu_state, color_image=None, depth_image=None):
        if not self.recording:
            return

        rgb_rel = ""
        depth_rel = ""

        if self.save_rgb and color_image is not None:
            rgb_name = f"{self.frame_id:06d}.jpg"
            rgb_path = self.rgb_dir / rgb_name

            ok = cv2.imwrite(
                str(rgb_path),
                color_image,
                [int(cv2.IMWRITE_JPEG_QUALITY), int(self.jpeg_quality)],
            )

            if ok:
                rgb_rel = "rgb/" + rgb_name

        if self.save_depth and depth_image is not None:
            depth_name = f"{self.frame_id:06d}.png"
            depth_path = self.depth_dir / depth_name

            ok = cv2.imwrite(str(depth_path), depth_image)

            if ok:
                depth_rel = "depth/" + depth_name

        row = {
            "frame_id": self.frame_id,

            "t_wall": time.time(),
            "t_monotonic": time.monotonic(),
            "t_esp_ms": packet["t_ms"],

            "din0_gpio25_record_switch": packet["din"][0],
            "din1_gpio26": packet["din"][1],
            "din2_gpio27_scissor_switch": packet["din"][2],
            "din3_gpio32": packet["din"][3],

            "dout0_gpio18_in1": packet["dout"][0],
            "dout1_gpio19_in2": packet["dout"][1],
            "scissor_state": infer_scissor_state(packet),

            "acc_x_g": packet["acc_g"][0],
            "acc_y_g": packet["acc_g"][1],
            "acc_z_g": packet["acc_g"][2],

            "gyro_x_dps": packet["gyro_dps"][0],
            "gyro_y_dps": packet["gyro_dps"][1],
            "gyro_z_dps": packet["gyro_dps"][2],

            "angle_x_deg": packet["angle_deg"][0],
            "angle_y_deg": packet["angle_deg"][1],
            "angle_z_deg": packet["angle_deg"][2],

            "temp_c": packet["temp_c"],

            "vel_x_ms": float(imu_state["velocity_ms"][0]),
            "vel_y_ms": float(imu_state["velocity_ms"][1]),
            "vel_z_ms": float(imu_state["velocity_ms"][2]),

            "lin_acc_x_ms2": float(imu_state["linear_acc_ms2"][0]),
            "lin_acc_y_ms2": float(imu_state["linear_acc_ms2"][1]),
            "lin_acc_z_ms2": float(imu_state["linear_acc_ms2"][2]),

            "acc_bias_x_ms2": float(imu_state["acc_bias_ms2"][0]),
            "acc_bias_y_ms2": float(imu_state["acc_bias_ms2"][1]),
            "acc_bias_z_ms2": float(imu_state["acc_bias_ms2"][2]),

            "bias_ready": int(imu_state["bias_ready"]),

            "rgb_path": rgb_rel,
            "depth_path": depth_rel,
        }

        self.writer.writerow(row)

        self.frame_id += 1

        if self.flush_every > 0 and self.frame_id % self.flush_every == 0:
            self.csv_file.flush()


# ============================================================
# Main recorder
# ============================================================

def run(args):
    if args.record_din == 2:
        print(
            "[WARNING] You selected din[2] / GPIO27 as recording trigger. "
            "But GPIO27 is already used as physical scissor switch."
        )

#     record_switch = SwitchChangeDetector(
#     debounce_s=args.debounce,
# )
#in case we need this former switch change detector, we can keep it here for reference. But the level controller is more suitable for this use case.
    record_switch = SwitchLevelController(
        active_value=args.record_active_value,
        debounce_s=args.debounce,
    )
    velocity_estimator = ProviderIMUVelocityEstimator(
    gravity_sign=args.gravity_sign,
    bias_seconds=args.bias_seconds,
    acc_deadband=args.acc_deadband,
    vel_decay=args.velocity_decay,
)

    writer = EpisodeWriter(
        root_dir=args.dataset,
        save_rgb=not args.no_rgb,
        save_depth=not args.no_depth,
        jpeg_quality=args.jpeg_quality,
        flush_every=args.flush_every,
    )

    camera = None

    try:
        if args.no_camera:
            print("[CAMERA] disabled by --no-camera.")
        else:
            camera = D405Camera(
                width=args.width,
                height=args.height,
                fps=args.fps,
                enable_depth=not args.no_depth,
                warmup_frames=args.camera_warmup,
            )
            camera.open()
            print("[CAMERA] D405 opened.")

        with ESP32Bridge(
            port=args.port,
            baud=args.baud,
            startup_wait=args.startup_wait,
            safe_stop_on_close=False,
        ) as bridge:

            print(f"[ESP32] Reading {args.port} at {args.baud} baud.")
            print(
                "[SWITCH] record trigger: "
                f"din[{args.record_din}], "
                f"mode=level-control, "
                f"active_value={args.record_active_value}, "
                f"debounce={args.debounce}s"
            )
            print(
                f"[INFO] DIN[{args.record_din}]={args.record_active_value} -> REC, "
                f"DIN[{args.record_din}]!={args.record_active_value} -> IDLE."
            )
            print("[INFO] Ctrl+C to exit.")

            packet_count = 0
            last_status_t = time.monotonic()

            while True:
                packet, drained = bridge.read_latest_packet()
                packet_count += 1 + drained
                serial_waiting = bridge.ser.in_waiting if bridge.ser is not None else -1
                raw_record_value = packet["din"][args.record_din]

                desired_recording = record_switch.update(raw_record_value)

                if desired_recording is not None:
                    if desired_recording and not writer.recording:
                        writer.start()
                        velocity_estimator.reset()

                    elif not desired_recording and writer.recording:
                        writer.stop()

                imu_state = velocity_estimator.update(packet)
                color_image = None
                depth_image = None

                if writer.recording and camera is not None:
                    color_image, depth_image = camera.read()

                writer.write(
                    packet=packet,
                    imu_state=imu_state,
                    color_image=color_image,
                    depth_image=depth_image,
                )

                now = time.monotonic()
                if now - last_status_t >= args.status_period:
                    last_status_t = now

                    rec_state = "REC" if writer.recording else "IDLE"

                    print(
                        f"[{rec_state}] "
                        f"packets={packet_count} "
                        f"DIN={packet['din']} "
                        f"DOUT={packet['dout']} "
                        f"scissor={infer_scissor_state(packet)} "
                        f"frames={writer.frame_id} "
                        f"drained={drained} "
                        f"serial_waiting={serial_waiting}"
                    )

    except KeyboardInterrupt:
        print("\n[EXIT] Interrupted by user.")

    finally:
        writer.stop()

        if camera is not None:
            camera.close()
            print("[CAMERA] closed.")


def parse_args():
    parser = argparse.ArgumentParser(
        description="Record ESP32 IMU/DIN/DOUT + D405 RGB/depth into episodes."
    )

    parser.add_argument(
        "--port",
        required=True,
        help="ESP32 serial port, for example /dev/ttyUSB0",
    )

    parser.add_argument(
        "--baud",
        type=int,
        default=115200,
        help="ESP32 USB serial baudrate.",
    )

    parser.add_argument(
        "--dataset",
        default="dataset",
        help="Output dataset root directory.",
    )

    parser.add_argument(
        "--record-din",
        type=int,
        default=0,
        choices=[0, 1, 2, 3],
        help=(
            "DIN index used as record start/stop switch. "
            "DIN_PINS=(25,26,27,32). "
            "0=GPIO25, 1=GPIO26, 2=GPIO27, 3=GPIO32. "
            "GPIO27/din2 is already used as scissor physical switch, "
            "so default is din0/GPIO25."
        ),
    )

    parser.add_argument(
        "--record-active-value",
        type=int,
        default=1,
        choices=[0, 1],
        help="Record when selected DIN equals this value. Default: 1, meaning DIN=1 records and DIN=0 stops.",
    )

    parser.add_argument(
        "--active-low",
        dest="active_low",
        action="store_true",
        default=True,
        help="Record switch is low-level active. Default: True.",
    )

    parser.add_argument(
        "--active-high",
        dest="active_low",
        action="store_false",
        help="Use high-level active record switch.",
    )

    parser.add_argument(
        "--debounce",
        type=float,
        default=0.25,
        help="Record switch debounce time in seconds.",
    )

    parser.add_argument(
        "--startup-wait",
        type=float,
        default=4.0,
        help="Wait time after opening ESP32 serial port.",
    )

    parser.add_argument(
        "--no-camera",
        action="store_true",
        help="Disable RealSense camera. Useful for testing DIN toggle and CSV only.",
    )

    parser.add_argument(
        "--width",
        type=int,
        default=640,
        help="RealSense color/depth width.",
    )

    parser.add_argument(
        "--height",
        type=int,
        default=480,
        help="RealSense color/depth height.",
    )

    parser.add_argument(
        "--fps",
        type=int,
        default=30,
        help="RealSense FPS.",
    )

    parser.add_argument(
        "--camera-warmup",
        type=int,
        default=15,
        help="Frames to discard after starting camera.",
    )

    parser.add_argument(
        "--no-rgb",
        action="store_true",
        help="Do not save RGB images.",
    )

    parser.add_argument(
        "--no-depth",
        action="store_true",
        help="Do not save depth images.",
    )

    parser.add_argument(
        "--jpeg-quality",
        type=int,
        default=95,
        help="JPEG quality for RGB image saving.",
    )

    parser.add_argument(
        "--velocity-decay",
        type=float,
        default=0.995,
        help="Decay factor for simple raw IMU velocity integration.",
    )

    parser.add_argument(
        "--flush-every",
        type=int,
        default=10,
        help="Flush CSV every N frames. Use 1 for safer but slower writing.",
    )

    parser.add_argument(
        "--status-period",
        type=float,
        default=1.0,
        help="Print status every N seconds.",
    )
    
    parser.add_argument(
    "--gravity-sign",
    type=float,
    default=1.0,
    choices=[1.0, -1.0],
    help="Use +1 if horizontal static acc_z is about +1g; use -1 if about -1g.",
)

    parser.add_argument(
        "--bias-seconds",
        type=float,
        default=1.0,
        help="Use the first N seconds after episode start to estimate residual acceleration bias.",
    )

    parser.add_argument(
        "--acc-deadband",
        type=float,
        default=0.10,
        help="Set linear acceleration to zero if norm is below this value, unit m/s^2.",
    )
    return parser.parse_args()


if __name__ == "__main__":
    run(parse_args())
# python3 record_episode.py --port /dev/ttyUSB0 --baud 115200 --record-din 0 --dataset dataset --no-camera