plot.py

# plot.py
# Plotting and visualization utilities for training/test curves and policy figures
# Also includes helpers to export Q-tables and test summaries (CSV/XLSX)

import os
import csv
from typing import List, Tuple, Dict, Optional

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

import config

# Optional Excel export via openpyxl
try:
    from openpyxl import Workbook
    from openpyxl.utils import get_column_letter
except Exception:
    Workbook = None
    get_column_letter = None


def ensure_dir(path: str) -> None:
    # Create directory if it doesn't exist
    os.makedirs(path, exist_ok=True)


def save_q_table(Q: np.ndarray, out_csv_path: str, out_npy_path: str) -> None:
    # Save Q-table in .csv and .npy formats for submission and debugging
    ensure_dir(os.path.dirname(out_csv_path) or ".")
    ensure_dir(os.path.dirname(out_npy_path) or ".")
    np.save(out_npy_path, Q)
    np.savetxt(out_csv_path, Q, delimiter=",")


# Publication-quality plotting style (used by comparison plots)
def _apply_pub_style() -> None:
    # Apply a consistent, publication-friendly style
    plt.rcParams.update({
        "figure.figsize": (6.6, 4.2),
        "figure.dpi": 120,
        "savefig.dpi": 600,
        "font.size": 11,
        "axes.titlesize": 12,
        "axes.labelsize": 11,
        "legend.fontsize": 10,
        "xtick.labelsize": 10,
        "ytick.labelsize": 10,
        "lines.linewidth": 2.0,
        "axes.linewidth": 1.0,
        "grid.linewidth": 0.8,
        "grid.alpha": 0.25,
        "pdf.fonttype": 42,   # Embed TrueType fonts in PDF
        "ps.fonttype": 42,
    })


def _save_pub_figure(fig: plt.Figure, out_path: str) -> None:
    # Save both PNG and PDF for report/paper figures
    ensure_dir(os.path.dirname(out_path) or ".")
    base, ext = os.path.splitext(out_path)
    if ext.lower() == "":
        png_path = base + ".png"
        pdf_path = base + ".pdf"
    else:
        png_path = base + ".png"
        pdf_path = base + ".pdf"

    fig.savefig(png_path, bbox_inches="tight", dpi=600)
    fig.savefig(pdf_path, bbox_inches="tight")
    plt.close(fig)


def _pretty_algo_name(algo: str) -> str:
    # Convert internal algo key to a display name in legends
    key = str(algo).strip().lower()
    if key in ("ql", "qlearning", "q-learning", "q_learning"):
        return "Q-learning"
    if key in ("sarsa",):
        return "SARSA"
    if key in ("mc", "montecarlo", "monte-carlo", "monte_carlo"):
        return "Monte Carlo"
    return str(algo)


def _order_algos_for_plot(keys: List[str]) -> List[str]:
    # Keep algorithm ordering consistent across figures
    preferred = ["mc", "sarsa", "ql"]
    keys_lower = {k: str(k).lower() for k in keys}
    ordered = []

    # Add preferred ones first
    for p in preferred:
        for k, kl in keys_lower.items():
            if kl == p:
                ordered.append(k)
                break

    # Add remaining (stable)
    for k in keys:
        if k not in ordered:
            ordered.append(k)

    return ordered


def _plot_multi_lines(
    series: Dict[str, Tuple[np.ndarray, np.ndarray]],
    title: str,
    xlabel: str,
    ylabel: str,
    out_path: str,
    use_markers: bool = False,
    legend_loc: str = "best",
) -> None:
    # Generic multi-line plot used for algorithm comparisons
    _apply_pub_style()
    fig, ax = plt.subplots()

    # Use different line styles/markers to distinguish algorithms
    linestyles = ["-", "--", "-.", ":"]
    markers = ["o", "s", "^", "D", "x"]

    ordered_keys = _order_algos_for_plot(list(series.keys()))

    for i, key in enumerate(ordered_keys):
        x, y = series[key]
        label = _pretty_algo_name(key)
        ls = linestyles[i % len(linestyles)]
        if use_markers:
            mk = markers[i % len(markers)]
            ax.plot(
                x, y,
                linestyle=ls,
                marker=mk,
                markersize=3.5,
                markevery=max(1, len(x) // 25),
                label=label,
            )
        else:
            ax.plot(x, y, linestyle=ls, label=label)

    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

    ax.grid(True)
    ax.legend(frameon=False, loc=legend_loc)

    # Remove top/right spines for cleaner plots
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    fig.tight_layout()
    _save_pub_figure(fig, out_path)


def plot_steps(steps: List[int], out_path: str) -> None:
    # Plot episode length during training
    ensure_dir(os.path.dirname(out_path) or ".")
    plt.figure()
    plt.plot(np.arange(len(steps)), steps)
    plt.title("")
    plt.xlabel("Episode")
    plt.ylabel("Steps")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()


def plot_average_rewards(avg_return: List[float], out_path: str) -> None:
    # Plot running average return during training
    ensure_dir(os.path.dirname(out_path) or ".")
    plt.figure()
    plt.plot(np.arange(len(avg_return)), avg_return)
    plt.title("")
    plt.xlabel("Episode")
    plt.ylabel("Average return")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()


def plot_accuracy(accuracy: List[float], window: int, out_path: str) -> None:
    # Plot windowed training success rate (recorded every window episodes)
    ensure_dir(os.path.dirname(out_path) or ".")
    plt.figure()
    x = np.arange(len(accuracy)) * int(window)
    plt.plot(x, accuracy)
    plt.title("")
    plt.xlabel("Episode")
    plt.ylabel("Success rate")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()


def plot_success_rate_from_accuracy(accuracy: List[float], window: int, out_path: str) -> None:
    # Alias used by main.py for compatibility
    plot_accuracy(accuracy, window, out_path)


def plot_test_steps(test_steps: List[int], out_path: str) -> None:
    # Plot test steps per episode during evaluation
    ensure_dir(os.path.dirname(out_path) or ".")
    plt.figure()
    plt.plot(np.arange(len(test_steps)), test_steps)
    plt.title("Test steps greedy")
    plt.xlabel("Test episode")
    plt.ylabel("Steps")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()


def plot_test_success_rate(test_success: List[int], out_path: str) -> None:
    # Plot 0/1 success result for each test episode
    ensure_dir(os.path.dirname(out_path) or ".")
    plt.figure()
    plt.plot(np.arange(len(test_success)), test_success)
    plt.title("Test success greedy")
    plt.xlabel("Test episode")
    plt.ylabel("Success 0/1")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()


def plot_steps_compare(all_steps: Dict[str, List[int]], out_path: str) -> None:
    # Compare training steps per episode across algorithms
    series: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
    for algo, steps in all_steps.items():
        x = np.arange(len(steps))
        y = np.array(steps, dtype=np.float64)
        series[algo] = (x, y)
    _plot_multi_lines(series, "", "Episode", "Steps", out_path)


def plot_avg_return_compare(all_avg: Dict[str, List[float]], out_path: str) -> None:
    # Compare training average return across algorithms
    series: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
    for algo, avg in all_avg.items():
        x = np.arange(len(avg))
        y = np.array(avg, dtype=np.float64)
        series[algo] = (x, y)
    _plot_multi_lines(series, "", "Episode", "Average return", out_path)


def plot_accuracy_compare(all_acc: Dict[str, List[float]], window: int, out_path: str) -> None:
    # Compare windowed success rate across algorithms
    series: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
    for algo, acc in all_acc.items():
        x = np.arange(len(acc)) * int(window)
        y = np.array(acc, dtype=np.float64)
        series[algo] = (x, y)

    # Legend location is fixed per project requirement
    _plot_multi_lines(
        series,
        "",
        "Episode",
        "Success rate",
        out_path,
        legend_loc="lower right",
    )


def plot_test_steps_compare(all_test_steps: Dict[str, List[int]], out_path: str) -> None:
    # Compare test steps across algorithms
    series: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
    for algo, steps in all_test_steps.items():
        x = np.arange(len(steps))
        y = np.array(steps, dtype=np.float64)
        series[algo] = (x, y)
    _plot_multi_lines(series, "", "Test episode", "Steps", out_path, use_markers=True)


def plot_test_success_compare(all_test_success: Dict[str, List[int]], out_path: str) -> None:
    # Compare test success 0/1 across algorithms
    series: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
    for algo, suc in all_test_success.items():
        x = np.arange(len(suc))
        y = np.array(suc, dtype=np.float64)
        series[algo] = (x, y)
    _plot_multi_lines(series, "", "Test episode", "Success 0/1", out_path, use_markers=True)


def _draw_grid_axes(ax, N: int) -> None:
    # Draw N x N grid with hidden ticks for a clean map figure
    ax.set_xlim(-0.5, N - 0.5)
    ax.set_ylim(N - 0.5, -0.5)  # invert y to match row-major indexing
    ax.set_aspect("equal")

    # Minor ticks define gridlines, but tick marks/labels are hidden
    ax.set_xticks(np.arange(-0.5, N, 1), minor=True)
    ax.set_yticks(np.arange(-0.5, N, 1), minor=True)
    ax.grid(which="minor", linewidth=0.8)

    ax.set_xticks([])
    ax.set_yticks([])
    ax.tick_params(which="both", bottom=False, left=False, labelbottom=False, labelleft=False)

    # Hide spines to avoid frame lines
    for spine in ax.spines.values():
        spine.set_visible(False)


def _action_to_vec(a: int) -> Tuple[float, float]:
    # Convert action id to a 2D arrow direction in plot coordinates
    # x is column direction, y is row direction
    if a == 0:
        return 0.0, -1.0
    if a == 1:
        return 0.0, 1.0
    if a == 2:
        return 1.0, 0.0
    if a == 3:
        return -1.0, 0.0
    return 0.0, 0.0


def save_best_action_map_image(env, Q: np.ndarray, out_path: str) -> None:
    # Save best-action map using arrows for all non-terminal states
    ensure_dir(os.path.dirname(out_path) or ".")
    N = int(env.grid_size)
    inner = env.map[1:-1, 1:-1]
    best_actions = np.argmax(Q, axis=1).reshape(N, N)

    # Build quiver inputs (skip traps and goal cells)
    X, Y, U, V = [], [], [], []
    for r in range(N):
        for c in range(N):
            cell = int(inner[r, c])
            if cell in (-1, 1):
                continue
            a = int(best_actions[r, c])
            dx, dy = _action_to_vec(a)
            X.append(c)
            Y.append(r)
            U.append(dx)
            V.append(dy)

    fig, ax = plt.subplots(figsize=(6.2, 6.2))
    _draw_grid_axes(ax, N)

    # Draw symbols for traps and goal
    for r in range(N):
        for c in range(N):
            if int(inner[r, c]) == -1:
                ax.text(c, r, "X", ha="center", va="center")
            elif int(inner[r, c]) == 1:
                ax.text(c, r, "G", ha="center", va="center")

    # Mark start position
    sr, sc = int(env.start_pos[0] - 1), int(env.start_pos[1] - 1)
    if 0 <= sr < N and 0 <= sc < N:
        ax.text(sc, sr, "S", ha="center", va="center")

    # Draw policy arrows
    if len(X) > 0:
        ax.quiver(
            np.array(X),
            np.array(Y),
            np.array(U),
            np.array(V),
            angles="xy",
            scale_units="xy",
            scale=2.5,
            width=0.006,
            alpha=0.85,
            pivot="mid",
        )

    ax.set_title("Best action map arrows")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close(fig)


def save_final_policy_image(env, Q: np.ndarray, out_path: str, max_steps: int) -> None:
    # Roll out greedy policy once and draw filled-cell path WITH best-action arrows overlaid
    # NOTE: Only light-blue filled cells are drawn for the path (no dark blue line).
    ensure_dir(os.path.dirname(out_path) or ".")
    N = int(env.grid_size)
    inner = env.map[1:-1, 1:-1]

    # Best actions for arrows (global policy visualization)
    best_actions = np.argmax(Q, axis=1).reshape(N, N)

    # Roll out one episode to get visited cells
    s = int(env.reset())
    pos = env.state_to_pos(s)
    r, c = int(pos[0] - 1), int(pos[1] - 1)

    path_cells = [(r, c)]
    for _ in range(int(max_steps)):
        a = int(np.argmax(Q[s]))
        res = env.step(a)
        s = int(res.next_state)

        pos = env.state_to_pos(s)
        r, c = int(pos[0] - 1), int(pos[1] - 1)
        path_cells.append((r, c))

        if res.done:
            break

    # Remove duplicates while preserving order for clean visualization
    seen = set()
    path_unique = []
    for rc in path_cells:
        if rc not in seen:
            path_unique.append(rc)
            seen.add(rc)

    # Prepare arrow vectors for each non-terminal cell
    X, Y, U, V = [], [], [], []
    for rr in range(N):
        for cc in range(N):
            cell = int(inner[rr, cc])
            if cell in (-1, 1):
                continue
            a = int(best_actions[rr, cc])
            dx, dy = _action_to_vec(a)
            X.append(cc)
            Y.append(rr)
            U.append(dx)
            V.append(dy)

    fig, ax = plt.subplots(figsize=(6.2, 6.2))
    _draw_grid_axes(ax, N)

    # Draw traps and goal markers
    for rr in range(N):
        for cc in range(N):
            if int(inner[rr, cc]) == -1:
                ax.text(cc, rr, "X", ha="center", va="center")
            elif int(inner[rr, cc]) == 1:
                ax.text(cc, rr, "G", ha="center", va="center")

    # Draw start marker
    sr, sc = int(env.start_pos[0] - 1), int(env.start_pos[1] - 1)
    if 0 <= sr < N and 0 <= sc < N:
        ax.text(sc, sr, "S", ha="center", va="center")

    # Draw arrows with lower opacity so the path shading stays visible
    if len(X) > 0:
        ax.quiver(
            np.array(X),
            np.array(Y),
            np.array(U),
            np.array(V),
            angles="xy",
            scale_units="xy",
            scale=2.7,
            width=0.0048,
            alpha=0.45,
            pivot="mid",
            zorder=2,
        )

    # Shade visited cells as full squares
    for (rr, cc) in path_unique:
        rect = patches.Rectangle(
            (cc - 0.5, rr - 0.5),
            1.0,
            1.0,
            linewidth=0.0,
            facecolor="lightskyblue",
            alpha=0.35,
            zorder=3,
        )
        ax.add_patch(rect)

    ax.set_title("")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close(fig)


def evaluate_greedy(env, Q: np.ndarray, num_episodes: int, max_steps: int) -> Tuple[List[int], List[float], List[int]]:
    # Run evaluation episodes using the greedy policy from Q
    # Returns per-episode steps, returns, and success flags
    test_steps: List[int] = []
    test_returns: List[float] = []
    test_success: List[int] = []

    for _ in range(int(num_episodes)):
        s = int(env.reset())
        ep_return = 0.0
        done_flag = False

        for t in range(int(max_steps)):
            a = int(np.argmax(Q[s]))
            res = env.step(a)
            ep_return += float(res.reward)
            s = int(res.next_state)

            if res.done:
                test_steps.append(t + 1)
                test_returns.append(ep_return)
                test_success.append(1 if float(res.reward) == float(config.REWARD_GOAL) else 0)
                done_flag = True
                break

        if not done_flag:
            test_steps.append(int(max_steps))
            test_returns.append(ep_return)
            test_success.append(0)

    return test_steps, test_returns, test_success


def summarize_test_results(test_steps: List[int], test_returns: List[float]) -> Tuple[float, float]:
    # Compute required summary stats for the report table
    steps_arr = np.array(test_steps, dtype=np.float64)
    rets_arr = np.array(test_returns, dtype=np.float64)
    var_test_steps = float(np.var(steps_arr)) if steps_arr.size > 0 else float("nan")
    avg_test_return = float(np.mean(rets_arr)) if rets_arr.size > 0 else float("nan")
    return var_test_steps, avg_test_return


def save_test_summary_xlsx(test_steps: List[int], test_returns: List[float], out_xlsx_path: str) -> None:
    # Save the 1-row summary table to an Excel file
    if Workbook is None:
        raise ImportError("openpyxl is not available. Please install it to save .xlsx files.")

    ensure_dir(os.path.dirname(out_xlsx_path) or ".")
    var_test_steps, avg_test_return = summarize_test_results(test_steps, test_returns)

    wb = Workbook()
    ws = wb.active
    ws.title = "test_summary"

    headers = ["var_test_steps", "avg_test_return"]
    values = [var_test_steps, avg_test_return]

    ws.append(headers)
    ws.append(values)

    # Adjust column width for readability in Excel
    if get_column_letter is not None:
        for col_idx, header in enumerate(headers, start=1):
            col_letter = get_column_letter(col_idx)
            ws.column_dimensions[col_letter].width = max(16, len(header) + 2)

    wb.save(out_xlsx_path)


def save_test_summary_csv(test_steps: List[int], test_returns: List[float], out_csv_path: str) -> None:
    # Save the 1-row summary table to CSV (Excel-friendly)
    ensure_dir(os.path.dirname(out_csv_path) or ".")
    var_test_steps, avg_test_return = summarize_test_results(test_steps, test_returns)

    with open(out_csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["var_test_steps", "avg_test_return"])
        writer.writerow([var_test_steps, avg_test_return])