import argparse
import csv
import os
from pathlib import Path

os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import matplotlib.pyplot as plt


TARGET_GROUPS = {
    "velocity": {
        "vel_x_ms": ["vel_x_ms"],
        "vel_y_ms": ["vel_y_ms"],
        "vel_z_ms": ["vel_z_ms"],
    },
    "lin_acc": {
        "lin_acc_x_ms2": ["lin_acc_x_ms2", "lin_acc_xms2"],
        "lin_acc_y_ms2": ["lin_acc_y_ms2", "lin_acc_yms2"],
        "lin_acc_z_ms2": ["lin_acc_z_ms2", "lin_acc_zms2"],
    },
}

TIME_CANDIDATES = ["t_monotonic", "t_wall", "t_esp_ms", "frame_id"]


def parse_args():
    parser = argparse.ArgumentParser(
        description="Read CSV files from a folder and plot velocity/linear acceleration signals."
    )
    parser.add_argument(
        "--input-dir",
        default="data_csv",
        help="Directory containing CSV files (default: data_csv)",
    )
    parser.add_argument(
        "--output-dir",
        default="png",
        help="Directory to save output images (default: png)",
    )
    parser.add_argument(
        "--dpi",
        type=int,
        default=140,
        help="Output image DPI (default: 140)",
    )
    return parser.parse_args()


def pick_first_existing(header, candidates):
    for name in candidates:
        if name in header:
            return name
    return None


def to_float(value):
    if value is None:
        return float("nan")
    text = str(value).strip()
    if text == "":
        return float("nan")
    try:
        return float(text)
    except ValueError:
        return float("nan")


def read_csv_data(csv_path):
    with csv_path.open("r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        header = reader.fieldnames or []

        time_col = pick_first_existing(header, TIME_CANDIDATES)
        resolved = {}
        for group in TARGET_GROUPS.values():
            for canonical_name, aliases in group.items():
                resolved[canonical_name] = pick_first_existing(header, aliases)

        rows = list(reader)

    if not rows:
        return None

    x = []
    if time_col is None:
        x = list(range(len(rows)))
        x_label = "sample_index"
    else:
        raw_t = [to_float(r.get(time_col)) for r in rows]
        t0 = raw_t[0] if raw_t else 0.0
        if time_col.endswith("_ms"):
            x = [((v - t0) / 1000.0) for v in raw_t]
            x_label = f"{time_col} (s, relative)"
        else:
            x = [(v - t0) for v in raw_t]
            x_label = f"{time_col} (relative)"

    data = {}
    for canonical_name, actual_name in resolved.items():
        if actual_name is None:
            data[canonical_name] = None
        else:
            data[canonical_name] = [to_float(r.get(actual_name)) for r in rows]

    return {
        "x": x,
        "x_label": x_label,
        "data": data,
        "resolved": resolved,
    }


def plot_one_csv(csv_path, output_dir, dpi):
    parsed = read_csv_data(csv_path)
    if parsed is None:
        print(f"[SKIP] {csv_path.name}: empty file")
        return False

    x = parsed["x"]
    x_label = parsed["x_label"]
    data = parsed["data"]
    resolved = parsed["resolved"]

    has_velocity = any(data[name] is not None for name in TARGET_GROUPS["velocity"].keys())
    has_lin_acc = any(data[name] is not None for name in TARGET_GROUPS["lin_acc"].keys())

    if not has_velocity and not has_lin_acc:
        print(f"[SKIP] {csv_path.name}: no target columns found")
        return False

    fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=True)
    fig.suptitle(f"Motion Signals - {csv_path.name}")

    ax_v, ax_a = axes

    for name in TARGET_GROUPS["velocity"].keys():
        series = data[name]
        if series is not None:
            ax_v.plot(x, series, label=name, linewidth=1.1)

    for name in TARGET_GROUPS["lin_acc"].keys():
        series = data[name]
        if series is not None:
            ax_a.plot(x, series, label=name, linewidth=1.1)

    ax_v.set_ylabel("velocity (m/s)")
    ax_a.set_ylabel("linear acc (m/s^2)")
    ax_a.set_xlabel(x_label)

    ax_v.grid(True, alpha=0.25)
    ax_a.grid(True, alpha=0.25)

    if has_velocity:
        ax_v.legend(loc="upper right")
    else:
        ax_v.text(0.5, 0.5, "No velocity columns", transform=ax_v.transAxes, ha="center", va="center")

    if has_lin_acc:
        ax_a.legend(loc="upper right")
    else:
        ax_a.text(0.5, 0.5, "No linear-acc columns", transform=ax_a.transAxes, ha="center", va="center")

    missing = [k for k, v in resolved.items() if v is None]
    if missing:
        fig.text(0.01, 0.01, f"Missing columns: {', '.join(missing)}", fontsize=9)

    fig.tight_layout(rect=[0, 0.03, 1, 0.97])

    output_dir.mkdir(parents=True, exist_ok=True)
    out_path = output_dir / f"{csv_path.stem}_motion.png"
    fig.savefig(out_path, dpi=dpi)
    plt.close(fig)

    print(f"[OK] {csv_path.name} -> {out_path}")
    return True


def main():
    args = parse_args()
    input_dir = Path(args.input_dir)
    output_dir = Path(args.output_dir)

    if not input_dir.exists() or not input_dir.is_dir():
        raise SystemExit(f"Input directory does not exist: {input_dir}")

    csv_files = sorted(input_dir.glob("*.csv"))
    if not csv_files:
        raise SystemExit(f"No CSV files found in: {input_dir}")

    ok_count = 0
    for csv_path in csv_files:
        if plot_one_csv(csv_path, output_dir, args.dpi):
            ok_count += 1

    print(f"Finished. Generated charts for {ok_count}/{len(csv_files)} files.")


if __name__ == "__main__":
    main()