"""Demonstrate runtime vibration-control toggling with persisted settings."""

import numpy as np
from pathlib import Path
from control.python.covalent_wrapper import (
    ShaperInterface,
    RobotState,
)
from control.python.control_settings import ControlSettingsManager

THIS_DIR = Path(__file__).resolve().parent
REPO_ROOT = THIS_DIR.parents[4]

# --- Configuration ---
Ts = 0.005  # 200Hz sampling
model_dir = REPO_ROOT / "src/robot/models/current"
num_axes = 3
num_joints = 6
python_src_root = REPO_ROOT / "src"
urdf_filepath = REPO_ROOT / "src/robot/urdf/test_robot.urdf"
settings = ControlSettingsManager()
print(f"Control settings path: {settings.config_path}")
print(f"Initial vibration control enabled: {settings.get_vibration_enabled()}")

# --- Initialize ---
shaper = ShaperInterface(
    sample_time=Ts,
    model_directory=str(model_dir),
    python_src_root=str(python_src_root),
    urdf_filepath=str(urdf_filepath),
    num_axes=num_axes,
    num_joints=num_joints,
)

# --- Example: shape a single position command vector ---
joint_positions = np.array([0.25, -0.3, 0.1, 0.0, 0.0, 0.0], dtype=float)
single_state = RobotState(joint_angles=joint_positions)

single_sample = shaper.process_sample(
    joint_positions,
    single_state,
    enable_vibration_control=settings.get_vibration_enabled(),
)
shaped_positions = single_sample.positions
print("Shaped joint positions:", shaped_positions)

# Simulate a runtime toggle event from UI/controller service.
settings.set_vibration_enabled(False)
print(
    f"After runtime toggle, vibration control enabled: {settings.get_vibration_enabled()}"
)
passthrough_sample = shaper.process_sample(
    joint_positions,
    single_state,
    enable_vibration_control=settings.get_vibration_enabled(),
)
print("Pass-through joint positions:", passthrough_sample.positions)

# --- Example: shape an entire trajectory (optional) ---
# Build a simple point-to-point move followed by a short dwell at the goal.
start_angles = np.array([0.0, -0.2, 0.1, 0.0, 0.0, 0.0], dtype=float)
goal_angles = np.array([0.4, 0.1, -0.2, 0.0, 0.0, 0.0], dtype=float)
move_duration_s = 1.0
dwell_duration_s = 0.3

move_samples = int(round(move_duration_s / Ts)) + 1
dwell_samples = int(round(dwell_duration_s / Ts))

alpha = np.linspace(0.0, 1.0, move_samples)
# Cubic smoothstep gives zero velocity at start/end of the move segment.
blend = 3.0 * alpha**2 - 2.0 * alpha**3
move_segment = start_angles + np.outer(blend, goal_angles - start_angles)
dwell_segment = np.repeat(goal_angles.reshape(1, -1), dwell_samples, axis=0)

trajectory = np.vstack([move_segment, dwell_segment])
N = trajectory.shape[0]
t = np.arange(N) * Ts

settings.set_vibration_enabled(True)
print(
    f"Re-enabled vibration control for batch example: {settings.get_vibration_enabled()}"
)
# Apply the shaper sample-by-sample or as a full
# Sample-by-sample stream
shaper_stream = ShaperInterface(
    sample_time=Ts,
    model_directory=str(model_dir),
    python_src_root=str(python_src_root),
    urdf_filepath=str(urdf_filepath),
    num_axes=num_axes,
    num_joints=num_joints,
)

shaped_series: list[tuple[float, np.ndarray]] = []
last_time = 0.0
# Developer note (Codex): read the toggle each cycle so UI changes are applied
# without restarting the control process.
for i in range(N):
    cmd = trajectory[i].copy()
    state = RobotState(joint_angles=cmd)
    sample = shaper_stream.process_sample(
        cmd,
        state,
        enable_vibration_control=settings.get_vibration_enabled(),
    )
    shaped_series.append((t[i], sample.positions))
    last_time = t[i]

tails = [shaper.finalize() for shaper in shaper_stream.shapers]
max_tail = max(len(tail) for tail in tails)
for tail_idx in range(max_tail):
    tail_time = last_time + (tail_idx + 1) * Ts
    pos = shaped_series[-1][1].copy()
    for axis in range(num_axes):
        if tail_idx < len(tails[axis]):
            y, v, a = tails[axis][tail_idx]
            pos[axis] = y
    shaped_series.append((tail_time, pos))

stream_len = len(shaped_series)

# With a trajectory shaper
shaper_batch = ShaperInterface(
    sample_time=Ts,
    model_directory=str(model_dir),
    python_src_root=str(python_src_root),
    urdf_filepath=str(urdf_filepath),
    num_axes=num_axes,
    num_joints=num_joints,
)

states = [RobotState(joint_angles=trajectory[i].copy()) for i in range(N)]
shaped_traj = shaper_batch.process_trajectory(
    trajectory,
    states,
    time_vector=list(t),
    enable_vibration_control=settings.get_vibration_enabled(),
)

batch_len = shaped_traj.positions.shape[0]
time_batch = np.arange(0, batch_len) * Ts

print(f"[StreamVsBatch] Streaming shaped samples (vibration enabled): {stream_len}")
print(f"[StreamVsBatch] Batch shaped samples (vibration enabled): {batch_len}")

# --- A/B test: vibration ON vs OFF with same input trajectory ---
shaper_enabled = ShaperInterface(
    sample_time=Ts,
    model_directory=str(model_dir),
    python_src_root=str(python_src_root),
    urdf_filepath=str(urdf_filepath),
    num_axes=num_axes,
    num_joints=num_joints,
)
shaper_disabled = ShaperInterface(
    sample_time=Ts,
    model_directory=str(model_dir),
    python_src_root=str(python_src_root),
    urdf_filepath=str(urdf_filepath),
    num_axes=num_axes,
    num_joints=num_joints,
)

traj_enabled = shaper_enabled.process_trajectory(
    trajectory,
    states,
    time_vector=list(t),
    enable_vibration_control=True,
)
traj_disabled = shaper_disabled.process_trajectory(
    trajectory,
    states,
    time_vector=list(t),
    enable_vibration_control=False,
)

enabled_samples = int(traj_enabled.positions.shape[0])
disabled_samples = int(traj_disabled.positions.shape[0])
enabled_end_time_s = float(traj_enabled.time[-1])
disabled_end_time_s = float(traj_disabled.time[-1])
tail_extension_s = enabled_end_time_s - disabled_end_time_s

common_samples = min(enabled_samples, disabled_samples)
max_abs_position_delta = float(
    np.max(
        np.abs(
            traj_enabled.positions[:common_samples, :num_axes]
            - traj_disabled.positions[:common_samples, :num_axes]
        )
    )
)
enabled_final_position = traj_enabled.positions[-1, :]
disabled_final_position = traj_disabled.positions[-1, :]
final_position_delta = enabled_final_position - disabled_final_position
max_abs_final_position_delta = float(np.max(np.abs(final_position_delta)))
final_position_tolerance_rad = 1e-6

print("\n[Toggle Comparison]")
print(
    f"Enabled: samples={enabled_samples}, end_time_s={enabled_end_time_s:.6f}, "
    f"Disabled: samples={disabled_samples}, end_time_s={disabled_end_time_s:.6f}"
)
print(f"Tail extension (enabled-disabled) [s]: {tail_extension_s:.6f}")
print(
    f"Max |position delta| over first {common_samples} samples on shaped axes [rad]: "
    f"{max_abs_position_delta:.6e}"
)
print(
    "Final position delta (enabled-disabled) [rad]: "
    f"{np.array2string(final_position_delta, precision=6)}"
)
print(f"Max |final position delta| [rad]: {max_abs_final_position_delta:.6e}")
if tail_extension_s > 0.0:
    print("PASS: Enabled trajectory includes shaper tail extension.")
else:
    print("WARN: No tail extension observed; investigate shaper configuration.")
if max_abs_final_position_delta <= final_position_tolerance_rad:
    print("PASS: Enabled and disabled trajectories end at the same final position.")
else:
    print(
        "WARN: Final positions differ beyond tolerance; "
        "check trajectory/state consistency."
    )
