import sys
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.dates as mdates  # Add this import at the top
import select
from datetime import datetime
import threading
import queue

fig = plt.figure(figsize=(10, 6))

ax_lr = fig.add_axes([0.05, 0.3, 0.9, 0.65])
ax_lr.set_facecolor("none")
ax_lr.set_ylim(bottom=0, top=None)
ax_lr.yaxis.set_ticks_position('right')
ax_lr.tick_params(axis='y', labelcolor='tab:brown')

ax_acc = fig.add_axes([0.05, 0.3, 0.9, 0.65], sharex=ax_lr)
ax_acc.set_facecolor("none")
ax_acc.set_ylim(0, 100)
ax_acc.yaxis.set_ticks_position('right')
ax_acc.tick_params(axis='y', labelcolor='tab:green')

ax_loss = fig.add_axes([0.05, 0.3, 0.9, 0.65], sharex=ax_lr)
ax_loss.set_facecolor("none")
ax_loss.set_ylim(bottom=0, top=None)
ax_loss.tick_params(axis='y', labelcolor='tab:blue')

ax_time = fig.add_axes([0.05, 0.05, 0.9, 0.2])
ax_time.set_facecolor("none")
ax_time.set_ylim(bottom=0, top=None)
ax_time.yaxis.set_ticks_position('right')
ax_time.tick_params(axis='y', labelcolor='tab:orange')

columns = ["index", "timestamp", "epoch", "batch", "lr", "loss", "accuracy", "grad_inp", "grad_qry", "grad_key", "grad_val", "grad_out", "queue_ms", "compute_ms", "infer_ms"]
df = pd.DataFrame(columns=columns)
ema_alpha = 0.1
window_size = 50

def parse_line(line):
    try:
        parts = line.strip().split(",")
        if parts[0] == "index" or len(parts) < len(columns):
            return None
        row = {
            "index": int(parts[0]),
            "timestamp": datetime.fromtimestamp(int(parts[1]) / 1000),
            "epoch": int(parts[2]),
            "batch": int(parts[3]),
            "lr": float(parts[4]),
            "loss": float(parts[5]),
            "accuracy": float(parts[6]) * 100,
            "grad_inp": float(parts[7]),
            "grad_qry": float(parts[8]),
            "grad_key": float(parts[9]),
            "grad_val": float(parts[10]),
            "grad_out": float(parts[11]),
            "queue_ms": float(parts[12]),
            "compute_ms": float(parts[13]),
            "infer_ms": float(parts[14]),
        }
        return row
    except Exception as e:
        print(f"Bad line: {line.strip()} ({e})", file=sys.stderr)
        return None

line_queue = queue.Queue()

def stdin_reader():
    for line in sys.stdin:
        line_queue.put(line)

threading.Thread(target=stdin_reader, daemon=True).start()

def get_new_lines():
    lines = []
    while not line_queue.empty():
        lines.append(line_queue.get())
    return lines

last_seen_index = None

def animate(_):
    global df

    lines = get_new_lines()
    if not lines:
        return

    new_rows = [parse_line(line) for line in lines]
    new_rows = [r for r in new_rows if r is not None]

    if not new_rows:
      return

    global last_seen_index

    new_index = new_rows[-1]["index"]
    if last_seen_index and (new_index < last_seen_index):
        print("Index reset, restarting plot")
        df = pd.DataFrame(columns=columns)

    last_seen_index = new_index

    df_new = pd.DataFrame(new_rows)
    df_new["total_ms"] = df_new["queue_ms"] + df_new["compute_ms"] + df_new["infer_ms"]
    df = pd.concat([df, df_new], ignore_index=True)

    df["ema_loss"] = df["loss"].ewm(alpha=ema_alpha).mean().combine_first(df["loss"])
    df["avg_loss"] = df["loss"].rolling(window=window_size).mean().combine_first(df["loss"])
    df["avg_accuracy"] = df["accuracy"].rolling(window=window_size).mean().combine_first(df["accuracy"])
    df["avg_total_ms"] = df["total_ms"].rolling(window=window_size).mean().combine_first(df["total_ms"])

    ax_lr.clear()
    ax_lr.set_ylabel("LR")
    ax_lr.yaxis.set_label_position("right")

    ax_acc.clear()
    ax_acc.set_ylabel("Accuracy (%)")

    ax_loss.clear()
    ax_loss.set_title("Live Training Metrics")
    ax_loss.set_xlabel("Epoch / Batch")
    ax_loss.set_ylabel("Loss")
    ax_acc.yaxis.set_label_position("right")

    ax_time.clear()
    ax_time.set_xlabel("Time")
    ax_time.set_ylabel("Processing Time (ms)")
    ax_time.yaxis.set_label_position("right")

    ax_lr.plot(df["index"], df["lr"], label="LR", linestyle='--', alpha=0.3, color="tab:brown")

    ax_acc.plot(df["index"], df["accuracy"], label="Accuracy", linestyle='--', alpha=0.3, color="tab:green")
    ax_acc.plot(df["index"], df["avg_accuracy"], label="Avg Accuracy",  color="tab:green")

    ax_loss.plot(df["index"], df["loss"], label="Loss", linestyle=':', alpha=0.3, color="tab:blue")
    ax_loss.plot(df["index"], df["ema_loss"], label="EMA Loss", linestyle='--', alpha=0.7, color="tab:blue")
    ax_loss.plot(df["index"], df["avg_loss"], label=f"Avg Loss", color="tab:blue")

    ax_time.plot(df["index"], df["queue_ms"], label="Queueing", linestyle=':', alpha=0.3, color="#ffa64d")
    ax_time.plot(df["index"], df["compute_ms"], label="Computing", linestyle=':', alpha=0.3, color="#cc6600")
    ax_time.plot(df["index"], df["infer_ms"], label="Inferring", linestyle=':', alpha=0.3, color="#ffb84d")
    ax_time.plot(df["index"], df["total_ms"], label="Total Time",linestyle='--', alpha=0.7, color="tab:orange")
    ax_time.plot(df["index"], df["avg_total_ms"], label="Avg Total Time", color="tab:orange")

    # Add epoch/batch label ticks to secondary X axis
    tick_locs = ax_time.get_xticks()
    tick_labels = []

    if not df.empty:
        for loc in tick_locs:
            idx = (abs(df["index"] - loc)).argmin()
            e, b = int(df.iloc[idx]["epoch"]), int(df.iloc[idx]["batch"])
            tick_labels.append(f"{e + 1}/{b + 1}")
    else:
        tick_labels = [""] * len(tick_locs)

    ax_loss.set_xlim(ax_time.get_xlim())
    ax_loss.set_xticks(tick_locs)
    ax_loss.set_xticklabels(tick_labels)

    # Gather all legend entries
    lines = []
    labels = []

    for ax in [ax_lr, ax_acc, ax_loss]:
        l, lab = ax.get_legend_handles_labels()
        lines += l
        labels += lab

    # Add combined legend to ax_loss
    ax_loss.legend(lines, labels, loc='upper left')  # or 'upper right', etc.
    ax_time.legend(loc='upper left')

ani = animation.FuncAnimation(fig, animate, interval=1000)
plt.show()