#!/usr/bin/env python3

import csv
import math
import sys
from datetime import datetime
from pathlib import Path
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile

from nav_msgs.msg import Odometry
from origincar_msg.msg import Data
from sensor_msgs.msg import Imu
from std_msgs.msg import Float32


def quaternion_to_yaw(x: float, y: float, z: float, w: float) -> float:
    siny_cosp = 2.0 * (w * z + x * y)
    cosy_cosp = 1.0 - 2.0 * (y * y + z * z)
    return math.atan2(siny_cosp, cosy_cosp)


class BaseFeedbackMonitor(Node):
    def __init__(self) -> None:
        super().__init__("base_feedback_monitor")
        qos = QoSProfile(depth=10)
        self.workspace_root = Path.cwd()
        self.log_dir = self.workspace_root / "feedback_logs" / datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.csv_path = self.log_dir / "feedback.csv"
        self.csv_file = self.csv_path.open("w", newline="", encoding="utf-8")
        self.csv_writer = csv.writer(self.csv_file)
        self.csv_writer.writerow([
            "wall_time",
            "elapsed_s",
            "imu_ready",
            "odom_ready",
            "voltage_ready",
            "robotpose_ready",
            "robotvel_ready",
            "gyrodebug_ready",
            "voltage_v",
            "imu_orientation_x",
            "imu_orientation_y",
            "imu_orientation_z",
            "imu_orientation_w",
            "imu_yaw_rad",
            "imu_angular_velocity_x",
            "imu_angular_velocity_y",
            "imu_angular_velocity_z",
            "imu_linear_acceleration_x",
            "imu_linear_acceleration_y",
            "imu_linear_acceleration_z",
            "odom_pose_x",
            "odom_pose_y",
            "odom_yaw_rad",
            "odom_twist_vx",
            "odom_twist_vy",
            "odom_twist_wz",
            "robotpose_x",
            "robotpose_y",
            "robotpose_z",
            "robotvel_x",
            "robotvel_y",
            "robotvel_z",
            "gyro_z_filtered_pre_bias",
            "gyro_z_bias_model",
            "gyro_z_final_for_yaw",
        ])

        self.start_time = self.get_clock().now()
        self.imu_msg: Optional[Imu] = None
        self.odom_msg: Optional[Odometry] = None
        self.voltage_msg: Optional[Float32] = None
        self.robotpose_msg: Optional[Data] = None
        self.robotvel_msg: Optional[Data] = None
        self.gyrodebug_msg: Optional[Data] = None

        self.create_subscription(Imu, "imu/data_raw", self.imu_callback, qos)
        self.create_subscription(Odometry, "odom", self.odom_callback, qos)
        self.create_subscription(Float32, "PowerVoltage", self.voltage_callback, qos)
        self.create_subscription(Data, "robotpose", self.robotpose_callback, qos)
        self.create_subscription(Data, "robotvel", self.robotvel_callback, qos)
        self.create_subscription(Data, "gyro_debug", self.gyrodebug_callback, qos)

        self.create_timer(0.2, self.render)
        self.create_timer(0.2, self.log_snapshot)
        self.get_logger().info(f"Logging feedback to {self.csv_path}")

    def imu_callback(self, msg: Imu) -> None:
        self.imu_msg = msg

    def odom_callback(self, msg: Odometry) -> None:
        self.odom_msg = msg

    def voltage_callback(self, msg: Float32) -> None:
        self.voltage_msg = msg

    def robotpose_callback(self, msg: Data) -> None:
        self.robotpose_msg = msg

    def robotvel_callback(self, msg: Data) -> None:
        self.robotvel_msg = msg

    def gyrodebug_callback(self, msg: Data) -> None:
        self.gyrodebug_msg = msg

    def _fmt_float(self, value: Optional[float], digits: int = 4) -> str:
        if value is None:
            return "N/A"
        return f"{value:.{digits}f}"

    def _topic_state(self, msg: object) -> str:
        return "OK" if msg is not None else "WAIT"

    def _get_snapshot(self) -> dict:
        imu_yaw = None
        odom_yaw = None

        if self.imu_msg is not None:
            imu_yaw = quaternion_to_yaw(
                self.imu_msg.orientation.x,
                self.imu_msg.orientation.y,
                self.imu_msg.orientation.z,
                self.imu_msg.orientation.w,
            )

        if self.odom_msg is not None:
            odom_yaw = quaternion_to_yaw(
                self.odom_msg.pose.pose.orientation.x,
                self.odom_msg.pose.pose.orientation.y,
                self.odom_msg.pose.pose.orientation.z,
                self.odom_msg.pose.pose.orientation.w,
            )

        return {
            "wall_time": datetime.now().isoformat(timespec="milliseconds"),
            "elapsed_s": (self.get_clock().now() - self.start_time).nanoseconds / 1e9,
            "imu_ready": int(self.imu_msg is not None),
            "odom_ready": int(self.odom_msg is not None),
            "voltage_ready": int(self.voltage_msg is not None),
            "robotpose_ready": int(self.robotpose_msg is not None),
            "robotvel_ready": int(self.robotvel_msg is not None),
            "gyrodebug_ready": int(self.gyrodebug_msg is not None),
            "voltage_v": self.voltage_msg.data if self.voltage_msg is not None else None,
            "imu_orientation_x": self.imu_msg.orientation.x if self.imu_msg is not None else None,
            "imu_orientation_y": self.imu_msg.orientation.y if self.imu_msg is not None else None,
            "imu_orientation_z": self.imu_msg.orientation.z if self.imu_msg is not None else None,
            "imu_orientation_w": self.imu_msg.orientation.w if self.imu_msg is not None else None,
            "imu_yaw_rad": imu_yaw,
            "imu_angular_velocity_x": self.imu_msg.angular_velocity.x if self.imu_msg is not None else None,
            "imu_angular_velocity_y": self.imu_msg.angular_velocity.y if self.imu_msg is not None else None,
            "imu_angular_velocity_z": self.imu_msg.angular_velocity.z if self.imu_msg is not None else None,
            "imu_linear_acceleration_x": self.imu_msg.linear_acceleration.x if self.imu_msg is not None else None,
            "imu_linear_acceleration_y": self.imu_msg.linear_acceleration.y if self.imu_msg is not None else None,
            "imu_linear_acceleration_z": self.imu_msg.linear_acceleration.z if self.imu_msg is not None else None,
            "odom_pose_x": self.odom_msg.pose.pose.position.x if self.odom_msg is not None else None,
            "odom_pose_y": self.odom_msg.pose.pose.position.y if self.odom_msg is not None else None,
            "odom_yaw_rad": odom_yaw,
            "odom_twist_vx": self.odom_msg.twist.twist.linear.x if self.odom_msg is not None else None,
            "odom_twist_vy": self.odom_msg.twist.twist.linear.y if self.odom_msg is not None else None,
            "odom_twist_wz": self.odom_msg.twist.twist.angular.z if self.odom_msg is not None else None,
            "robotpose_x": self.robotpose_msg.x if self.robotpose_msg is not None else None,
            "robotpose_y": self.robotpose_msg.y if self.robotpose_msg is not None else None,
            "robotpose_z": self.robotpose_msg.z if self.robotpose_msg is not None else None,
            "robotvel_x": self.robotvel_msg.x if self.robotvel_msg is not None else None,
            "robotvel_y": self.robotvel_msg.y if self.robotvel_msg is not None else None,
            "robotvel_z": self.robotvel_msg.z if self.robotvel_msg is not None else None,
            "gyro_z_filtered_pre_bias": self.gyrodebug_msg.x if self.gyrodebug_msg is not None else None,
            "gyro_z_bias_model": self.gyrodebug_msg.y if self.gyrodebug_msg is not None else None,
            "gyro_z_final_for_yaw": self.gyrodebug_msg.z if self.gyrodebug_msg is not None else None,
        }

    def log_snapshot(self) -> None:
        snapshot = self._get_snapshot()
        self.csv_writer.writerow(snapshot.values())
        self.csv_file.flush()

    def render(self) -> None:
        snapshot = self._get_snapshot()
        sys.stdout.write("\033[2J\033[H")
        sys.stdout.write("origincar base feedback monitor\n")
        sys.stdout.write("=" * 40 + "\n")
        sys.stdout.write(
            f"imu/data_raw: {self._topic_state(self.imu_msg)} | "
            f"odom: {self._topic_state(self.odom_msg)} | "
            f"PowerVoltage: {self._topic_state(self.voltage_msg)} | "
            f"robotpose: {self._topic_state(self.robotpose_msg)} | "
            f"robotvel: {self._topic_state(self.robotvel_msg)} | "
            f"gyro_debug: {self._topic_state(self.gyrodebug_msg)}\n\n"
        )
        sys.stdout.write(f"log file: {self.csv_path}\n\n")

        if self.voltage_msg is not None:
            sys.stdout.write(f"voltage: {self._fmt_float(self.voltage_msg.data, 3)} V\n")
        else:
            sys.stdout.write("voltage: N/A\n")

        if self.imu_msg is not None:
            imu = self.imu_msg
            sys.stdout.write(
                "imu orientation: "
                f"x={self._fmt_float(imu.orientation.x)} "
                f"y={self._fmt_float(imu.orientation.y)} "
                f"z={self._fmt_float(imu.orientation.z)} "
                f"w={self._fmt_float(imu.orientation.w)} "
                f"yaw={self._fmt_float(snapshot['imu_yaw_rad'])} rad\n"
            )
            sys.stdout.write(
                "imu angular vel: "
                f"x={self._fmt_float(imu.angular_velocity.x)} "
                f"y={self._fmt_float(imu.angular_velocity.y)} "
                f"z={self._fmt_float(imu.angular_velocity.z)} rad/s\n"
            )
            sys.stdout.write(
                "imu linear acc:  "
                f"x={self._fmt_float(imu.linear_acceleration.x)} "
                f"y={self._fmt_float(imu.linear_acceleration.y)} "
                f"z={self._fmt_float(imu.linear_acceleration.z)} m/s^2\n"
            )
        else:
            sys.stdout.write("imu orientation: N/A\n")
            sys.stdout.write("imu angular vel: N/A\n")
            sys.stdout.write("imu linear acc:  N/A\n")

        if self.odom_msg is not None:
            odom = self.odom_msg
            sys.stdout.write(
                "odom pose:       "
                f"x={self._fmt_float(odom.pose.pose.position.x)} "
                f"y={self._fmt_float(odom.pose.pose.position.y)} "
                f"yaw={self._fmt_float(snapshot['odom_yaw_rad'])}\n"
            )
            sys.stdout.write(
                "odom twist:      "
                f"vx={self._fmt_float(odom.twist.twist.linear.x)} "
                f"vy={self._fmt_float(odom.twist.twist.linear.y)} "
                f"wz={self._fmt_float(odom.twist.twist.angular.z)}\n"
            )
        else:
            sys.stdout.write("odom pose:       N/A\n")
            sys.stdout.write("odom twist:      N/A\n")

        if self.robotpose_msg is not None:
            pose = self.robotpose_msg
            sys.stdout.write(
                "robotpose:       "
                f"x={self._fmt_float(pose.x)} "
                f"y={self._fmt_float(pose.y)} "
                f"z={self._fmt_float(pose.z)}\n"
            )
        else:
            sys.stdout.write("robotpose:       N/A\n")

        if self.robotvel_msg is not None:
            vel = self.robotvel_msg
            sys.stdout.write(
                "robotvel:        "
                f"x={self._fmt_float(vel.x)} "
                f"y={self._fmt_float(vel.y)} "
                f"z={self._fmt_float(vel.z)}\n"
            )
        else:
            sys.stdout.write("robotvel:        N/A\n")

        if self.gyrodebug_msg is not None:
            gyro = self.gyrodebug_msg
            sys.stdout.write(
                "gyro debug:      "
                f"pre_bias={self._fmt_float(gyro.x, 6)} "
                f"bias_z={self._fmt_float(gyro.y, 6)} "
                f"final_z={self._fmt_float(gyro.z, 6)} rad/s\n"
            )
        else:
            sys.stdout.write("gyro debug:      N/A\n")

        sys.stdout.write("\nCtrl+C to exit\n")
        sys.stdout.flush()


def main(args=None) -> None:
    rclpy.init(args=args)
    node = BaseFeedbackMonitor()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.csv_file.close()
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
