"""
Closed-loop gripper force control from two tactile sensors.

Run from repository root:
    python examples/gripper_force_control.py --port COM3 -v1 0 -v2 1

The controller reads left/right tactile normal force and updates the gripper
target force percentage. It uses the larger side as the safety signal because
the current gripper SDK exposes one shared force command.
"""

import argparse
import os
import sys
import time
from pathlib import Path

import multiprocessing as mp

from two_sensor_force_reader import TwoSensorForceReader


REPO_ROOT = Path(__file__).resolve().parents[1]
SDK_DIR = REPO_ROOT / "sdk"
if str(SDK_DIR) not in sys.path:
    sys.path.insert(0, str(SDK_DIR))


class ExponentialFilter:
    def __init__(self, alpha, initial=0.0):
        self.alpha = max(0.0, min(1.0, float(alpha)))
        self.value = float(initial)
        self.initialized = False

    def update(self, sample):
        sample = float(sample)
        if not self.initialized:
            self.value = sample
            self.initialized = True
        else:
            self.value = self.alpha * sample + (1.0 - self.alpha) * self.value
        return self.value


def clamp(value, min_value, max_value):
    return max(min_value, min(max_value, value))


def step_toward_open(current_pos, open_pos, delta):
    if current_pos is None:
        return open_pos
    if current_pos > open_pos:
        return max(open_pos, current_pos - delta)
    if current_pos < open_pos:
        return min(open_pos, current_pos + delta)
    return open_pos


def calibrate_baseline(reader, seconds):
    if seconds <= 0:
        return 0.0, 0.0

    left_values = []
    right_values = []
    deadline = time.perf_counter() + seconds
    while time.perf_counter() < deadline:
        left, right = reader.get_forces(timeout=0.2)
        if left is not None:
            left_values.append(left)
        if right is not None:
            right_values.append(right)
        time.sleep(0.01)

    left_baseline = sum(left_values) / len(left_values) if left_values else 0.0
    right_baseline = sum(right_values) / len(right_values) if right_values else 0.0
    return left_baseline, right_baseline


def compute_next_force_pct(
    current_pct,
    control_force,
    target_force,
    deadband,
    force_min,
    force_max,
    step_up,
    step_down,
):
    low = target_force - deadband
    high = target_force + deadband

    if control_force > high:
        return clamp(current_pct - step_down, force_min, force_max), "down"
    if control_force < low:
        return clamp(current_pct + step_up, force_min, force_max), "up"
    return current_pct, "hold"


def load_gripper(args):
    from changingtek_p_rtu_Servo import MotorController

    return MotorController(
        args.port,
        args.slave_id,
        baudrate=args.baudrate,
        timeout=args.serial_timeout,
    )


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument("--video1", "-v1", type=str, default="0", help="left sensor camera index or video path")
    parser.add_argument("--video2", "-v2", type=str, default="1", help="right sensor camera index or video path")
    parser.add_argument("--rotation1", "-r1", type=str, default="./config/rotation_config_0.json", help="left sensor orientation config")
    parser.add_argument("--rotation2", "-r2", type=str, default="./config/rotation_config_1.json", help="right sensor orientation config")
    parser.add_argument("--sensor-config1", type=str, default="./config/ddjx01.json", help="left tactile sensor config")
    parser.add_argument("--sensor-config2", type=str, default="./config/ddjx01.json", help="right tactile sensor config")
    parser.add_argument("--sensor-fps", type=float, default=30.0, help="target processing FPS per tactile sensor")
    parser.add_argument("--motion-threshold", type=float, default=0.0, help="motion threshold used by tactile deformation")
    parser.add_argument("--sensor-cpu", action="store_true", help="disable CUDA for tactile workers")

    parser.add_argument("--port", type=str, default="COM5", help="gripper serial port")
    parser.add_argument("--slave-id", type=int, default=1, help="Modbus slave id")
    parser.add_argument("--baudrate", type=int, default=115200, help="serial baudrate")
    parser.add_argument("--serial-timeout", type=float, default=0.2, help="serial timeout in seconds")

    parser.add_argument("--open-pos", type=int, default=0, help="open gripper position")
    parser.add_argument("--close-pos", type=int, default=9000, help="close gripper position")
    parser.add_argument("--speed", type=int, default=5, help="gripper speed percent")
    parser.add_argument("--accel", type=int, default=60, help="gripper acceleration")
    parser.add_argument("--decel", type=int, default=60, help="gripper deceleration")

    parser.add_argument("--initial-force", type=int, default=10, help="initial gripper force percent")
    parser.add_argument("--force-min", type=int, default=5, help="minimum gripper force percent")
    parser.add_argument("--force-max", type=int, default=35, help="maximum gripper force percent")
    parser.add_argument("--step-up", type=int, default=1, help="force percent increase per control tick")
    parser.add_argument("--step-down", type=int, default=2, help="force percent decrease per control tick")
    parser.add_argument("--control-hz", type=float, default=10.0, help="gripper force update frequency")

    parser.add_argument("--target-touch-force", type=float, default=800.0, help="target tactile normal force value")
    parser.add_argument("--deadband", type=float, default=80.0, help="do not adjust inside target +/- deadband")
    parser.add_argument("--emergency-touch-force", type=float, default=1600.0, help="open slightly if either side exceeds this value")
    parser.add_argument("--emergency-open-delta", type=int, default=500, help="position step toward open on emergency")
    parser.add_argument("--imbalance-threshold", type=float, default=0.0, help="optional left/right force difference warning threshold; 0 disables")
    parser.add_argument("--filter-alpha", type=float, default=0.25, help="force low-pass filter alpha")
    parser.add_argument("--baseline-seconds", type=float, default=1.0, help="seconds to sample no-contact tactile baseline")

    parser.add_argument("--no-initial-move", action="store_true", help="do not send initial close command")
    parser.add_argument("--trigger-force-update", action="store_true", help="call trigger_motion after set_target_force")
    parser.add_argument("--dry-run", action="store_true", help="print commands without opening the serial port")

    return parser.parse_args()


def main():
    args = parse_args()

    reader = TwoSensorForceReader(
        video1=args.video1,
        video2=args.video2,
        config1=args.sensor_config1,
        config2=args.sensor_config2,
        rotation1=args.rotation1,
        rotation2=args.rotation2,
        target_fps=args.sensor_fps,
        cuda=not args.sensor_cpu,
        motion_threshold=args.motion_threshold,
    )

    gripper = None
    force_pct = int(clamp(args.initial_force, args.force_min, args.force_max))
    left_filter = ExponentialFilter(args.filter_alpha)
    right_filter = ExponentialFilter(args.filter_alpha)
    last_written_force = None

    print("Starting tactile sensors...")
    reader.start()

    try:
        print(f"Calibrating tactile baseline for {args.baseline_seconds:.2f}s. Keep sensors unloaded.")
        left_base, right_base = calibrate_baseline(reader, args.baseline_seconds)
        print(f"baseline: left={left_base:.4f}, right={right_base:.4f}")

        if not args.dry_run:
            gripper = load_gripper(args)

        if not args.no_initial_move:
            if args.dry_run:
                print(
                    "[dry-run] temp_move "
                    f"position={args.close_pos}, speed={args.speed}, force={force_pct}, "
                    f"accel={args.accel}, decel={args.decel}"
                )
            else:
                gripper.temp_move(
                    position_mm=args.close_pos,
                    speed_pct=args.speed,
                    force_pct=force_pct,
                    accel=args.accel,
                    decel=args.decel,
                    trigger=True,
                )
            last_written_force = force_pct

        interval = 1.0 / args.control_hz if args.control_hz > 0 else 0.1
        print("Closed-loop control started. Press Ctrl+C to stop.")

        while True:
            tick_start = time.perf_counter()
            left_raw, right_raw = reader.get_forces(timeout=1.0)
            if left_raw is None or right_raw is None:
                print(f"waiting for forces: left={left_raw}, right={right_raw}")
                time.sleep(0.05)
                continue

            left_force = max(0.0, float(left_raw) - left_base)
            right_force = max(0.0, float(right_raw) - right_base)
            left_filtered = left_filter.update(left_force)
            right_filtered = right_filter.update(right_force)

            control_force = max(left_filtered, right_filtered)
            avg_force = 0.5 * (left_filtered + right_filtered)
            imbalance = left_filtered - right_filtered

            if control_force >= args.emergency_touch_force:
                force_pct = args.force_min
                action = "emergency"
                if args.dry_run:
                    print(f"[dry-run] set_target_force {force_pct}")
                    print(f"[dry-run] open toward {args.open_pos} by {args.emergency_open_delta}")
                else:
                    gripper.set_target_force(force_pct)
                    try:
                        current_pos = gripper.read_real_position()
                    except Exception:
                        current_pos = None
                    open_pos = step_toward_open(
                        current_pos,
                        args.open_pos,
                        args.emergency_open_delta,
                    )
                    gripper.temp_move(
                        position_mm=int(open_pos),
                        speed_pct=args.speed,
                        force_pct=force_pct,
                        accel=args.accel,
                        decel=args.decel,
                        trigger=True,
                    )
                last_written_force = force_pct
            else:
                force_pct, action = compute_next_force_pct(
                    force_pct,
                    control_force,
                    args.target_touch_force,
                    args.deadband,
                    args.force_min,
                    args.force_max,
                    args.step_up,
                    args.step_down,
                )

                if force_pct != last_written_force:
                    if args.dry_run:
                        print(f"[dry-run] set_target_force {force_pct}")
                    else:
                        gripper.set_target_force(int(force_pct))
                        if args.trigger_force_update:
                            gripper.trigger_motion()
                    last_written_force = force_pct

            imbalance_msg = ""
            if args.imbalance_threshold > 0 and abs(imbalance) >= args.imbalance_threshold:
                imbalance_msg = " imbalance"

            print(
                f"L={left_filtered:8.2f} R={right_filtered:8.2f} "
                f"avg={avg_force:8.2f} max={control_force:8.2f} "
                f"diff={imbalance:8.2f} force_pct={force_pct:3d} "
                f"action={action}{imbalance_msg}"
            )

            elapsed = time.perf_counter() - tick_start
            sleep_time = interval - elapsed
            if sleep_time > 0:
                time.sleep(sleep_time)

    except KeyboardInterrupt:
        print("\nStopping closed-loop control.")
    finally:
        reader.stop()


if __name__ == "__main__":
    mp.freeze_support()
    main()
