import argparse
import ast
import csv
import os
import re
import sys
import time
from datetime import datetime
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import matplotlib.pyplot as plt
import numpy as np
import serial

from pc_reader import PACKET_SIZE, SYNC, decode_packet


TEXT_LINE_RE = re.compile(
    r"#?\s*(?P<count>\d+)\s+"
    r"(?P<freq>[-+]?\d+(?:\.\d+)?)Hz\s+"
    r"DIN=(?P<din>\[[^\]]+\])\s+"
    r"DOUT=(?P<dout>\[[^\]]+\])\s+"
    r"ACC=(?P<acc>\[[^\]]+\])\s+"
    r"GYRO=(?P<gyro>\[[^\]]+\])\s+"
    r"ANGLE=(?P<angle>\[[^\]]+\])"
)


FIELD_NAMES = (
    "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",
    "freq_hz",
)


def read_serial_samples(port, baud, seconds=None, max_samples=None, max_gap_ms=1000):
    samples = []
    buffer = bytearray()
    last_t_ms = None
    elapsed_s = 0.0
    bad_delta_count = 0

    ser = serial.Serial()
    ser.port = port
    ser.baudrate = baud
    ser.timeout = 0.05
    ser.write_timeout = 0.05

    # 关键：避免 ESP32 因 DTR/RTS 被 pyserial 拉动而自动 reset / 进入 boot mode
    ser.dtr = False
    ser.rts = False

    ser.open()

    try:
        ser.dtr = False
        ser.rts = False
        ser.reset_input_buffer()
        ser.reset_output_buffer()

        print("Reading {} at {} baud...".format(port, baud), file=sys.stderr)
        print("Waiting for ESP32 main.py startup...", file=sys.stderr)
        time.sleep(4.0)

        start = time.monotonic()

        while True:
            if seconds is not None and time.monotonic() - start >= seconds:
                break
            if max_samples is not None and len(samples) >= max_samples:
                break

            chunk = ser.read(ser.in_waiting or 1)
            if chunk:
                buffer.extend(chunk)

            while len(buffer) >= PACKET_SIZE:
                sync_index = buffer.find(SYNC)
                if sync_index < 0:
                    del buffer[:-1]
                    break
                if sync_index:
                    del buffer[:sync_index]
                if len(buffer) < PACKET_SIZE:
                    break

                packet = bytes(buffer[:PACKET_SIZE])
                if (sum(packet[:-1]) & 0xFF) != packet[-1]:
                    del buffer[0]
                    continue

                del buffer[:PACKET_SIZE]
                data = decode_packet(packet)

                if last_t_ms is None:
                    freq_hz = 0.0
                else:
                    delta = (data["t_ms"] - last_t_ms) & 0xFFFFFFFF
                    if 0 < delta <= max_gap_ms:
                        elapsed_s += delta / 1000.0
                        freq_hz = 1000.0 / delta
                    else:
                        bad_delta_count += 1
                        freq_hz = 0.0

                last_t_ms = data["t_ms"]
                samples.append(flatten_packet(data, elapsed_s, freq_hz))

                if max_samples is not None and len(samples) >= max_samples:
                    break

    finally:
        ser.close()

    if bad_delta_count:
        print(
            "Ignored {} abnormal packet time delta(s).".format(bad_delta_count),
            file=sys.stderr,
        )

    return samples


def flatten_packet(data, t_s, freq_hz):
    return {
        "t_s": t_s,
        "acc_x_g": data["acc_g"][0],
        "acc_y_g": data["acc_g"][1],
        "acc_z_g": data["acc_g"][2],
        "gyro_x_dps": data["gyro_dps"][0],
        "gyro_y_dps": data["gyro_dps"][1],
        "gyro_z_dps": data["gyro_dps"][2],
        "angle_x_deg": data["angle_deg"][0],
        "angle_y_deg": data["angle_deg"][1],
        "angle_z_deg": data["angle_deg"][2],
        "temp_c": data["temp_c"],
        "freq_hz": freq_hz,
        "din0": data["din"][0],
        "din1": data["din"][1],
        "din2": data["din"][2],
        "din3": data["din"][3],
        "dout0": data["dout"][0],
        "dout1": data["dout"][1],
    }


def read_text_log(path):
    samples = []
    t_s = 0.0

    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            match = TEXT_LINE_RE.search(line)
            if not match:
                continue

            freq_hz = float(match.group("freq"))
            din = ast.literal_eval(match.group("din"))
            dout = ast.literal_eval(match.group("dout"))
            acc = ast.literal_eval(match.group("acc"))
            gyro = ast.literal_eval(match.group("gyro"))
            angle = ast.literal_eval(match.group("angle"))

            if samples and freq_hz > 0:
                t_s += 1.0 / freq_hz

            samples.append(
                {
                    "t_s": t_s,
                    "acc_x_g": acc[0],
                    "acc_y_g": acc[1],
                    "acc_z_g": acc[2],
                    "gyro_x_dps": gyro[0],
                    "gyro_y_dps": gyro[1],
                    "gyro_z_dps": gyro[2],
                    "angle_x_deg": angle[0],
                    "angle_y_deg": angle[1],
                    "angle_z_deg": angle[2],
                    "temp_c": np.nan,
                    "freq_hz": freq_hz,
                    "din0": din[0],
                    "din1": din[1],
                    "din2": din[2],
                    "din3": din[3],
                    "dout0": dout[0],
                    "dout1": dout[1],
                }
            )

    return samples


def write_csv(samples, path):
    if not samples:
        return

    columns = list(samples[0].keys())
    with open(path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=columns)
        writer.writeheader()
        writer.writerows(samples)

def euler_to_rot_matrix(roll_deg, pitch_deg, yaw_deg):
    """
    Convert IMU Euler angles to body->world rotation matrix.
    Assumption:
        angle_x = roll
        angle_y = pitch
        angle_z = yaw
    Rotation order:
        R = Rz(yaw) @ Ry(pitch) @ Rx(roll)

    注意：这个顺序要和 IWT603 官方定义确认。
    如果发现补偿后静止时 linear_acc 不接近 0.需要调整欧拉角顺序或正负号。
    """
    roll = np.deg2rad(roll_deg)
    pitch = np.deg2rad(pitch_deg)
    yaw = np.deg2rad(yaw_deg)

    cr, sr = np.cos(roll), np.sin(roll)
    cp, sp = np.cos(pitch), np.sin(pitch)
    cy, sy = np.cos(yaw), np.sin(yaw)

    Rx = np.array([
        [1, 0, 0],
        [0, cr, -sr],
        [0, sr, cr],
    ])

    Ry = np.array([
        [cp, 0, sp],
        [0, 1, 0],
        [-sp, 0, cp],
    ])

    Rz = np.array([
        [cy, -sy, 0],
        [sy, cy, 0],
        [0, 0, 1],
    ])

    return Rz @ Ry @ Rx


def add_linear_velocity(samples, gravity_sign=1.0, acc_deadband=0.15, vel_decay=0.995, bias_seconds=1.0):
    if not samples:
        return samples

    # 先计算每一帧的重力补偿线加速度
    raw_linear_acc_list = []

    for sample in samples:
        acc_body_g = np.array([
            sample["acc_x_g"],
            sample["acc_y_g"],
            sample["acc_z_g"],
        ], dtype=float)

        R = euler_to_rot_matrix(
            sample["angle_x_deg"],
            sample["angle_y_deg"],
            sample["angle_z_deg"],
        )

        acc_world_g = R @ acc_body_g
        linear_acc_world_g = acc_world_g - np.array([0.0, 0.0, gravity_sign])
        linear_acc_world = linear_acc_world_g * 9.81

        raw_linear_acc_list.append(linear_acc_world)

    raw_linear_acc_arr = np.array(raw_linear_acc_list)

    # 用前 bias_seconds 秒估计静止零偏
    t0 = samples[0]["t_s"]
    bias_indices = [
        i for i, s in enumerate(samples)
        if s["t_s"] - t0 <= bias_seconds
    ]

    if bias_indices:
        acc_bias = np.mean(raw_linear_acc_arr[bias_indices], axis=0)
    else:
        acc_bias = np.zeros(3)

    print("Estimated linear acc bias m/s^2:", acc_bias)

    v = np.zeros(3, dtype=float)
    last_t = samples[0]["t_s"]

    for i, sample in enumerate(samples):
        t = sample["t_s"]
        dt = t - last_t
        last_t = t

        if dt <= 0 or dt > 0.2:
            dt = 0.0

        linear_acc_world = raw_linear_acc_arr[i] - acc_bias

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

        v = v * vel_decay + linear_acc_world * dt

        sample["lin_acc_x_ms2"] = linear_acc_world[0]
        sample["lin_acc_y_ms2"] = linear_acc_world[1]
        sample["lin_acc_z_ms2"] = linear_acc_world[2]
        sample["vel_x_ms"] = v[0]
        sample["vel_y_ms"] = v[1]
        sample["vel_z_ms"] = v[2]

    return samples
def unique_path(path):
    """
    Automatically save CSV into ./csv/
    Automatically save PNG into ./png/
    Create folders if they do not exist.
    """

    if not path:
        return path

    root, ext = os.path.splitext(os.path.basename(path))
    ext = ext.lower()

    if ext == ".csv":
        folder = "csv"
    elif ext == ".png":
        folder = "png"
    else:
        folder = "output"

    os.makedirs(folder, exist_ok=True)

    new_path = os.path.join(folder, root + ext)

    if not os.path.exists(new_path):
        return new_path

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    return os.path.join(folder, "{}_{}{}".format(root, timestamp, ext))

def print_stats(samples):
    if not samples:
        print("No samples parsed.")
        return

    print("Samples: {}".format(len(samples)))
    print("Duration: {:.3f}s".format(samples[-1]["t_s"] - samples[0]["t_s"]))

    for group_name, fields in (
        ("ACC g", ("acc_x_g", "acc_y_g", "acc_z_g")),
        ("GYRO dps", ("gyro_x_dps", "gyro_y_dps", "gyro_z_dps")),
        ("ANGLE deg", ("angle_x_deg", "angle_y_deg", "angle_z_deg")),
    ):
        print("\n{}".format(group_name))
        for field in fields:
            values = np.array([sample[field] for sample in samples], dtype=float)
            diffs = np.diff(values)
            diff_std = float(np.std(diffs)) if len(diffs) else 0.0
            print(
                "  {:12s} mean={: .6f} std={:.6f} ptp={:.6f} diff_std={:.6f}".format(
                    field,
                    float(np.mean(values)),
                    float(np.std(values)),
                    float(np.ptp(values)),
                    diff_std,
                )
            )

    freq = np.array([sample["freq_hz"] for sample in samples[1:]], dtype=float)
    if len(freq):
        print(
            "\nFrequency Hz mean={:.3f} std={:.3f} min={:.3f} max={:.3f}".format(
                float(np.mean(freq)),
                float(np.std(freq)),
                float(np.min(freq)),
                float(np.max(freq)),
            )
        )
    if "lin_acc_x_ms2" in samples[0]:
        print("\nLINEAR ACC m/s^2")
        for field in ("lin_acc_x_ms2", "lin_acc_y_ms2", "lin_acc_z_ms2"):
            values = np.array([sample[field] for sample in samples], dtype=float)
            diffs = np.diff(values)
            diff_std = float(np.std(diffs)) if len(diffs) else 0.0
            print(
                "  {:16s} mean={: .6f} std={:.6f} ptp={:.6f} diff_std={:.6f}".format(
                    field,
                    float(np.mean(values)),
                    float(np.std(values)),
                    float(np.ptp(values)),
                    diff_std,
                )
            )

    if "vel_x_ms" in samples[0]:
        print("\nVELOCITY m/s")
        for field in ("vel_x_ms", "vel_y_ms", "vel_z_ms"):
            values = np.array([sample[field] for sample in samples], dtype=float)
            print(
                "  {:16s} mean={: .6f} std={:.6f} ptp={:.6f} final={: .6f}".format(
                    field,
                    float(np.mean(values)),
                    float(np.std(values)),
                    float(np.ptp(values)),
                    float(values[-1]),
                )
            )


def plot_samples(samples, output=None, show=True):
    if not samples:
        raise ValueError("No samples to plot.")

    t = np.array([sample["t_s"] for sample in samples], dtype=float)
    fig, axes = plt.subplots(6, 1, sharex=True, figsize=(13, 13))

    plot_group(
        axes[0],
        t,
        samples,
        ("acc_x_g", "acc_y_g", "acc_z_g"),
        "Raw Acc (g)",
    )

    plot_group(
        axes[1],
        t,
        samples,
        ("lin_acc_x_ms2", "lin_acc_y_ms2", "lin_acc_z_ms2"),
        "Linear Acc (m/s²)",
    )

    plot_group(
        axes[2],
        t,
        samples,
        ("vel_x_ms", "vel_y_ms", "vel_z_ms"),
        "Velocity (m/s)",
    )

    plot_group(
        axes[3],
        t,
        samples,
        ("gyro_x_dps", "gyro_y_dps", "gyro_z_dps"),
        "Gyro (deg/s)",
    )

    plot_group(
        axes[4],
        t,
        samples,
        ("angle_x_deg", "angle_y_deg", "angle_z_deg"),
        "Angle (deg)",
    )

    freq = np.array([sample["freq_hz"] for sample in samples], dtype=float)
    axes[5].plot(t, freq, label="freq_hz", linewidth=1.0)
    axes[5].set_ylabel("Hz")
    axes[5].set_xlabel("Time (s)")
    axes[5].grid(True, alpha=0.3)
    axes[5].legend(loc="upper right")

    fig.suptitle("IMU Data Quality Overview")
    fig.tight_layout()

    if output:
        fig.savefig(output, dpi=160)
        print("Saved plot to {}".format(output))

    if show:
        plt.show()


def plot_group(axis, t, samples, fields, ylabel):
    for field in fields:
        values = np.array([sample[field] for sample in samples], dtype=float)
        axis.plot(t, values, label=field, linewidth=1.0)
    axis.set_ylabel(ylabel)
    axis.grid(True, alpha=0.3)
    axis.legend(loc="upper right", ncol=3)


def main():
    parser = argparse.ArgumentParser(
        description="Read ESP IMU data and visualise signal quality/noise."
    )
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--file", help="pc_reader text log, for example rcd.txt")
    source.add_argument("--port", help="ESP32 serial port, for example /dev/ttyUSB0")
    parser.add_argument("--baud", type=int, default=115200)
    parser.add_argument("--seconds", type=float, default=10.0)
    parser.add_argument("--samples", type=int)
    parser.add_argument(
        "--max-gap-ms",
        type=int,
        default=1000,
        help="Ignore packet timestamp jumps larger than this.",
    )
    parser.add_argument("--output", default="imu_quality.png")
    parser.add_argument("--csv", help="Optional CSV export path.")
    parser.add_argument("--no-show", action="store_true", help="Save only; do not open a window.")
    args = parser.parse_args()

    if args.file:
        samples = read_text_log(args.file)
    else:
        samples = read_serial_samples(
            args.port,
            args.baud,
            args.seconds,
            args.samples,
            args.max_gap_ms,
        )

    # 关键：没有样本就直接退出，避免后面 CSV/plot 误报或报错
    if not samples:
        print("No samples parsed. Check input mode, serial port, baudrate, or log file format.")
        return

    samples = add_linear_velocity(samples, gravity_sign=1.0)
    print_stats(samples)

    if args.csv:
        csv_path = unique_path(args.csv)
        write_csv(samples, csv_path)
        print("Saved CSV to {}".format(csv_path))

    output_path = unique_path(args.output)
    plot_samples(samples, output=output_path, show=not args.no_show)

if __name__ == "__main__":
    main()
# python3 visualise.py --port /dev/ttyUSB0 --baud 115200 --seconds 10 --csv imu_data.csv