import os
import re
import sys
from collections import defaultdict

import matplotlib

try:
    matplotlib.use("QtAgg")
except Exception:
    matplotlib.use("Agg")

import matplotlib.pyplot as plt

# ================= INPUT =================
if len(sys.argv) < 2:
    print("Usage: python analyze_log.py <logfile> [output_dir]")
    sys.exit(1)

logfile = sys.argv[1]
out_dir = sys.argv[2] if len(sys.argv) >= 3 else "out"

# ================= PATH LOGIC =================
base_name = os.path.splitext(os.path.basename(logfile))[0]

# normalize folder structure
out_dir = os.path.normpath(out_dir)

pics_dir = os.path.join(out_dir, "pics")
tables_dir = os.path.join(out_dir, "tables")

os.makedirs(pics_dir, exist_ok=True)
os.makedirs(tables_dir, exist_ok=True)

output_png = os.path.join(pics_dir, f"{base_name}.png")
output_md = os.path.join(tables_dir, f"{base_name}.md")


# ================= PARSE =================
pattern = re.compile(r"(START|END).*TID=(\d+).*range=\[(\d+),(\d+)\].*time=([\d.]+)")

events = defaultdict(dict)

with open(logfile) as f:
    for line in f:
        m = pattern.search(line)
        if not m:
            continue

        typ, tid, l, r, t = m.groups()
        key = (tid, int(l), int(r))
        events[key][typ] = float(t)


# ================= BUILD ROWS =================
rows = []

for (tid, l, r), v in events.items():
    if "START" in v and "END" in v:
        start = v["START"]
        end = v["END"]
        duration = end - start

        rows.append(
            {
                "tid": tid,
                "range": f"[{l},{r}]",
                "start": start,
                "end": end,
                "duration": duration,
            }
        )

rows.sort(key=lambda x: x["start"])

if not rows:
    print("No valid events found")
    sys.exit(1)


# ================= OFFSET =================
t0 = rows[0]["start"]
for r in rows:
    r["offset"] = r["start"] - t0


# ================= COLOR MAP (LOGICAL) =================
unique_tids = sorted(set(r["tid"] for r in rows))

color_map = {
    tid: plt.cm.tab20(i % 20)  # stable, readable palette
    for i, tid in enumerate(unique_tids)
}


# ================= SAVE MARKDOWN =================
with open(output_md, "w") as f:
    f.write("# Execution Table\n\n")
    f.write("| TID | Range | Start | End | Duration | Offset |\n")
    f.write("|-----|-------|-------|-----|----------|--------|\n")

    for r in rows:
        f.write(
            f"| {r['tid']} | {r['range']} | "
            f"{r['start']:.6f} | {r['end']:.6f} | "
            f"{r['duration']:.6f} | {r['offset']:.6f} |\n"
        )

print(f"[OK] Table saved: {output_md}")


# ================= PLOT =================
plt.figure(figsize=(12, 6))

for i, r in enumerate(rows):
    plt.plot(
        [r["offset"], r["offset"] + r["duration"]],
        [i, i],
        color=color_map[r["tid"]],
        linewidth=4,
    )

# legend (TID → color)
for tid in unique_tids:
    plt.plot([], [], color=color_map[tid], label=f"TID {tid}")

plt.legend(loc="upper right", fontsize=8)

plt.xlabel("Time (seconds from start)")
plt.ylabel("Tasks")
plt.title(f"Execution Timeline: {base_name}")
plt.grid(True)
plt.tight_layout()


# ================= SAVE IMAGE =================
plt.savefig(output_png, dpi=200)
print(f"[OK] Plot saved: {output_png}")


# ================= SHOW =================
# plt.show()