import os
import numpy as np
from vigs.util.compute_recall import compute_recall_from_file
from scripts.align_ply_from_ape_log import align_ply_from_ape_log

# ── user config ────────────────────────────────────────────────────────────────
DATA_ROOT          = "data/VIGS-SLAM"
output_folder      = "outputs/output_vigs/opensource"
config_file        = "config/vigs.yaml"
traj_name          = "traj_kf_beforeBA.txt"  # or "traj_kf_afterBA.txt" for keyframe trajectory after final BA, "traj_full_beforeBA.txt" for full trajectory before final BA, "traj_full_afterBA.txt" for full trajectory after final BA
stride             = 1
IMU_poseinit_after = 20
recall_thresh_cm   = 10

SEQS = [
    "Basement1", "Basement2", "Basement3", "Basement4",
    "Downtown1", "Downtown2", "Downtown3",
    "Ferrari1", "Ferrari2",
    "Graffiti1", "Graffiti2", "Graffiti3",
    "Motorworld",
    "NTU_Campus1", "NTU_Campus2",
    "NTU_Corridor1", "NTU_Corridor2",
    "NTU_Office",
]

# Per-sequence buffer sizes (keyframe capacity). Keeping these tight reduces GPU memory usage (making it able to run on GPU with less memory) — set each value slightly above the actual keyframe count for the sequence.
BUFFER_SIZES = {
    "Basement1": 321, "Basement2": 439, "Basement3": 389, "Basement4": 311,
    "Downtown1": 180, "Downtown2": 389, "Downtown3": 943,
    "Ferrari1": 1180, "Ferrari2": 476,
    "Graffiti1": 186, "Graffiti2": 339, "Graffiti3": 412,
    "Motorworld": 399,
    "NTU_Campus1": 1479, "NTU_Campus2": 2862,
    "NTU_Corridor1": 2107, "NTU_Corridor2": 1680,
    "NTU_Office": 463,
}
# ───────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    os.makedirs(output_folder, exist_ok=True)

    ate_values, scale_errors, recalls = [], [], []

    for name in SEQS:
        seq = os.path.join(DATA_ROOT, name)
        out = os.path.join(output_folder, name)
        os.makedirs(out, exist_ok=True)
        print(f"{'='*60}\n{name}\n{'='*60}")

        buffer = BUFFER_SIZES[name] + 50
        cmd = (f"python demo.py"
               f" --calib {seq}/calib.txt"
               f" --imagedir {seq}/images"
               f" --config {config_file}"
               f" --stride {stride}"
               f" --IMU_poseinit_after {IMU_poseinit_after}"
               f" --imufile {seq}/imu.txt"
               f" --output {out}"
               f" --buffer {buffer}"
               f" --undistort"
               + f" > {out}/log.txt")
        print("RUNNING COMMAND: ", cmd)
        if not os.path.exists(f"{out}/{traj_name}"):
            os.system(cmd)
            # sbatch alternative:
            # os.system(f"sbatch --time=4:00:00 -n 1 --cpus-per-task=8 --mem-per-cpu=7G --gpus=1 --gres=gpumem:20g --output=slurms/{name}.out --wrap '{cmd}'")

        gt_file   = f"{seq}/MT-Pose.txt"
        traj_stem = traj_name.split(".")[0]
        log_ape   = f"{out}/log_ape_{traj_stem}.txt"
        os.system(f"evo_ape tum -vas --no_warnings --plot_mode xy"
                  f" --save_plot {out}/ape_se3_{traj_stem}.png"
                  f" --save_results {out}/ape_results.zip"
                  f" {gt_file} {out}/{traj_name}"
                  f" > {log_ape}")

        try:
            lines = open(log_ape).readlines()
            ATE   = float([l for l in lines if "rmse"             in l][-1].split("\t")[-1])
            scale = float([l for l in lines if "Scale correction" in l][-1].split(" ")[-1])
        except Exception as e:
            print(f"  [WARN] could not parse APE log: {e}")
            continue

        ATE_cm        = ATE * 100
        scale_err_pct = abs(1 - scale) * 100
        _, __, recall = compute_recall_from_file(gt_file, f"{out}/{traj_name}", thresh_cm=recall_thresh_cm)

        ate_values.append(ATE_cm)
        scale_errors.append(scale_err_pct)
        recalls.append(recall)
        print(f"  ATE={ATE_cm:.2f} cm  scale_err={scale_err_pct:.2f}%  recall@{recall_thresh_cm}cm={recall:.2f}%")

        try:
            if "beforeBA" in traj_name:
                align_ply_from_ape_log(ape_log_path=log_ape,
                                       ply_in=f"{out}/3dgs_before_final.ply",
                                       ply_out=f"{out}/3dgs_before_final_aligned.ply")
            elif "afterBA" in traj_name:
                align_ply_from_ape_log(ape_log_path=log_ape,
                                       ply_in=f"{out}/3dgs_final.ply",
                                       ply_out=f"{out}/3dgs_final_aligned.ply")
        except Exception as e:
            print(f"  [WARN] PLY alignment failed: {e}")

    if ate_values:
        print("\n--- CSV ---")
        print("," + ",".join(f"{v:.2f}" for v in ate_values))
        print("," + ",".join(f"{v:.2f}" for v in scale_errors))
        print("," + ",".join(f"{v:.2f}" for v in recalls))
        print("\n--- LaTeX ---")
        print("ATE [cm] & "          + " & ".join(f"{v:.2f}" for v in ate_values)   + f" & {np.mean(ate_values):.2f} \\\\")
        print("Scale error [\\%] & " + " & ".join(f"{v:.2f}" for v in scale_errors) + f" & {np.mean(scale_errors):.2f} \\\\")
        print(f"Recall@{recall_thresh_cm}cm [\\%] & " + " & ".join(f"{v:.2f}" for v in recalls) + f" & {np.mean(recalls):.2f} \\\\")
