plot.py

#!/usr/bin/env python3
"""
plot_clean.py

Plot CA grid output from DOTABLE(grid.txt) without loading the whole file.
"""

from __future__ import annotations

import math
from pathlib import Path
from typing import Dict, Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap, PowerNorm


MODE_PRESETS = {
    "interactive": {
        "font.size": 14,
        "axes.titlesize": 16,
        "axes.labelsize": 15,
        "xtick.labelsize": 13,
        "ytick.labelsize": 13,
        "legend.fontsize": 13,
        "figure.titlesize": 16,
        "seam_linewidth": 1.5,
        "save_dpi": 160,
        "figsize": (8.4, 10.2),
    },
    "publication": {
        "font.size": 16,
        "axes.titlesize": 18,
        "axes.labelsize": 17,
        "xtick.labelsize": 15,
        "ytick.labelsize": 15,
        "legend.fontsize": 15,
        "figure.titlesize": 18,
        "seam_linewidth": 1.8,
        "save_dpi": 300,
        "figsize": (8.8, 10.6),
    },
}


def ask_mode() -> str:
#    ans = input("Mode: interactive or publication? [interactive]: ").strip().lower()
#    if ans in ("", "i", "interactive", "screen"):
#        return "interactive"
#    if ans in ("p", "pub", "publication", "print"):
#        return "publication"
#    print("Unrecognized mode; using interactive.")
    return "publication"


def apply_mode(mode: str) -> dict:
    preset = MODE_PRESETS[mode]
    plt.rcParams.update({
        "font.size": preset["font.size"],
        "axes.titlesize": preset["axes.titlesize"],
        "axes.labelsize": preset["axes.labelsize"],
        "xtick.labelsize": preset["xtick.labelsize"],
        "ytick.labelsize": preset["ytick.labelsize"],
        "legend.fontsize": preset["legend.fontsize"],
        "figure.titlesize": preset["figure.titlesize"],
    })
    return preset


def ask_save_dpi(default_dpi: int) -> Tuple[bool, int, Optional[str]]:
    ans = input("Save figure to file? [y/N]: ").strip().lower()
    if ans not in ("y", "yes"):
        return False, default_dpi, None

    dpi_ans = input(f"Save DPI [{default_dpi}]: ").strip()
    if dpi_ans == "":
        dpi = default_dpi
    else:
        dpi = max(72, int(dpi_ans))

    out_ans = input("Output image filename [auto]: ").strip()
    if out_ans == "":
        out_name = None
    else:
        out_name = out_ans
    return True, dpi, out_name


def build_output_name(grid_path: str, mode: str, tile2x: bool) -> str:
    stem = Path(grid_path).stem
    suffix = "tile2x" if tile2x else "grid"
    return f"{stem}_{mode}_{suffix}.png"


def build_cmap() -> LinearSegmentedColormap:
    colors = [
        "#fff27a",
        "#d9f04a",
        "#8eea4d",
        "#39d96b",
        "#17cfa8",
        "#10c8e8",
        "#1698ff",
    ]
    cmap = LinearSegmentedColormap.from_list("ca_ppm_style_pop", colors, N=256)
    cmap.set_bad(color="black")
    return cmap


def parse_line(line: str) -> Optional[Tuple[int, int, float, float, float, float, float, float]]:
    s = line.rstrip("\n")
    if not s.strip():
        return None
    if s.lstrip().startswith("X Y"):
        return None

    parts = s.split()
    if len(parts) >= 8:
        try:
            return (
                int(parts[0]),
                int(parts[1]),
                float(parts[2]),
                float(parts[3]),
                float(parts[4]),
                float(parts[5]),
                float(parts[6]),
                float(parts[7]),
            )
        except ValueError:
            pass

    try:
        return (
            int(s[0:4]),
            int(s[4:8]),
            float(s[8:18]),
            float(s[18:28]),
            float(s[28:38]),
            float(s[38:48]),
            float(s[48:58]),
            float(s[58:68]),
        )
    except (ValueError, IndexError):
        return None


def scan_bounds_and_temps(path: str) -> Tuple[int, int, int, int, Dict[int, float]]:
    xmin = math.inf
    xmax = -math.inf
    ymin = math.inf
    ymax = -math.inf
    temp_by_y: Dict[int, float] = {}

    with open(path, "r", encoding="utf-8", errors="replace") as f:
        for line in f:
            parsed = parse_line(line)
            if parsed is None:
                continue
            x, y, _, _, _, _, temp, _ = parsed
            xmin = min(xmin, x)
            xmax = max(xmax, x)
            ymin = min(ymin, y)
            ymax = max(ymax, y)
            if y not in temp_by_y:
                temp_by_y[y] = temp

    if not math.isfinite(xmin) or not math.isfinite(ymin):
        raise ValueError("No valid grid data found in file.")

    return int(xmin), int(xmax), int(ymin), int(ymax), temp_by_y


def ask_range(name: str, lo: int, hi: int) -> Tuple[int, int]:
    ans = input(f"Use full {name} range [{lo}, {hi}]? [Y/n]: ").strip().lower()
    if ans in ("", "y", "yes"):
        return lo, hi

    v0 = int(input(f"{name} min: ").strip())
    v1 = int(input(f"{name} max: ").strip())
    if v0 > v1:
        v0, v1 = v1, v0
    v0 = max(lo, v0)
    v1 = min(hi, v1)
    return v0, v1


def ask_stride(name: str) -> int:
    ans = input(f"{name} stride [1]: ").strip()
    if ans == "":
        return 1
    stride = int(ans)
    return max(1, stride)


def fill_subset(
    path: str,
    xmin: int,
    xmax: int,
    ymin: int,
    ymax: int,
    xstride: int,
    ystride: int,
) -> np.ndarray:
    nx = 1 + (xmax - xmin) // xstride
    ny = 1 + (ymax - ymin) // ystride
    comp = np.full((ny, nx), np.nan, dtype=np.float32)

    with open(path, "r", encoding="utf-8", errors="replace") as f:
        for line in f:
            parsed = parse_line(line)
            if parsed is None:
                continue
            x, y, liq, lcomp, _, _, _, _ = parsed

            if x < xmin or x > xmax or y < ymin or y > ymax:
                continue
            if (x - xmin) % xstride != 0 or (y - ymin) % ystride != 0:
                continue

            ix = (x - xmin) // xstride
            iy = (y - ymin) // ystride

            if np.isfinite(liq) and liq > 0.0 and np.isfinite(lcomp):
                comp[iy, ix] = lcomp

    return comp


def compute_limits(comp: np.ndarray) -> Tuple[float, float]:
    vals = comp[np.isfinite(comp)]
    if vals.size == 0:
        return 0.0, 1.0

    vmin = float(np.percentile(vals, 2.0))
    vmax = float(np.percentile(vals, 98.0))

    if vmax <= vmin:
        vmin = float(np.min(vals))
        vmax = float(np.max(vals))
    if vmax <= vmin:
        vmax = vmin + 1.0
    return vmin, vmax


def build_temp_axis(
    ymin: int,
    ymax: int,
    ystride: int,
    temp_by_y: Dict[int, float],
) -> Tuple[np.ndarray, np.ndarray]:
    ys = sorted(y for y in temp_by_y if ymin <= y <= ymax)
    if not ys:
        return np.array([], dtype=float), np.array([], dtype=float)

    ypix = np.array([(y - ymin) / ystride + 1.0 for y in ys], dtype=float)
    temps = np.array([temp_by_y[y] for y in ys], dtype=float)
    return ypix, temps


def main() -> None:
    mode = ask_mode()
    preset = apply_mode(mode)

    path = input("Enter the path to the grid output text file: ").strip()
    xmin0, xmax0, ymin0, ymax0, temp_by_y = scan_bounds_and_temps(path)

    print(f"Detected grid bounds: X = [{xmin0}, {xmax0}], Y = [{ymin0}, {ymax0}]")

    xmin, xmax = ask_range("X", xmin0, xmax0)
    ymin, ymax = ask_range("Y", ymin0, ymax0)
    xstride = ask_stride("X")
    ystride = ask_stride("Y")

    tile2x = input("Tile X by 2 for periodic display? [Y/n]: ").strip().lower()
    tile2x = tile2x in ("", "y", "yes")

    comp = fill_subset(path, xmin, xmax, ymin, ymax, xstride, ystride)
    ny, nx = comp.shape
    comp_plot = np.hstack((comp, comp)) if tile2x else comp
    nx_plot = comp_plot.shape[1]

    cmap = build_cmap()
    vmin, vmax = compute_limits(comp)
    norm = PowerNorm(gamma=0.65, vmin=vmin, vmax=vmax)

    fig, ax = plt.subplots(figsize=preset["figsize"])
    fig.subplots_adjust(left=0.20, right=0.88)

    x0 = 0.5
    x1 = nx_plot + 0.5
    y0 = 0.5
    y1 = ny + 0.5

    ax.set_box_aspect(ny / max(nx_plot, 1))
    im = ax.imshow(
        comp_plot,
        cmap=cmap,
        norm=norm,
        origin="lower",
        extent=[x0, x1, y0, y1],
        interpolation="none",
    )

    if tile2x:
        ax.axvline(nx + 0.5, linewidth=preset["seam_linewidth"])

    ax.set_xlim(0.0, float(nx_plot))
    ax.set_ylim(0.0, float(ny))
    ax.set_aspect("equal", adjustable="box")
    ax.autoscale(enable=False)

    xtick_step = 100 if nx_plot >= 400 else (50 if nx_plot >= 50 else max(1, nx_plot // 5 or 1))
    ytick_step = 50 if ny >= 50 else max(1, ny // 5 or 1)

    if tile2x:
        xticks_left = np.arange(0, nx + 1, xtick_step)
        xticks_right = np.arange(nx, nx_plot + 1, xtick_step)
        xticks = np.unique(np.concatenate((xticks_left, xticks_right)))
        if xticks[-1] != nx_plot:
            xticks = np.append(xticks, nx_plot)
    else:
        xticks = np.arange(0, nx_plot + 1, xtick_step)
        if xticks[-1] != nx_plot:
            xticks = np.append(xticks, nx_plot)

    yticks = np.arange(0, ny + 1, ytick_step)
    if yticks[-1] != ny:
        yticks = np.append(yticks, ny)

    ax.set_xticks(xticks)
    ax.set_yticks(yticks)

    def x_label_from_tick(t: float) -> str:
        if tile2x:
            if t <= nx:
                return str((xmin - 1) + int(t * xstride))
            return str((xmin - 1) + int((t - nx) * xstride))
        return str((xmin - 1) + int(t * xstride))

    ax.set_xticklabels([x_label_from_tick(t) for t in xticks])
    ax.set_yticklabels([str((ymin - 1) + int(t * ystride)) for t in yticks])

    xlabel = "X (cells)"
    if tile2x:
        xlabel += " — periodic (tiled 2×)"
    ax.set_xlabel(xlabel)
    ax.set_ylabel("Y (cells)")

    cax = fig.add_axes([0.06, 0.15, 0.025, 0.70])
    cbar = fig.colorbar(im, cax=cax)
    cbar.set_label("Liquid composition Lcomp (wt%)")
    cbar.ax.tick_params(labelsize=max(11, preset["xtick.labelsize"] - 1))

    ax2 = ax.twinx()
    ax2.set_ylim(ax.get_ylim())
    ax2.set_yticks(yticks)
    ax2.set_ylabel("Temperature (°C)")
    ax2.tick_params(labelsize=preset["ytick.labelsize"])

    ypix, temps = build_temp_axis(ymin, ymax, ystride, temp_by_y)
    if len(ypix) >= 2 and np.isfinite(temps).all():
        ax2.set_yticklabels([f"{np.interp(t, ypix, temps):.1f}" for t in yticks])
    elif len(ypix) == 1 and np.isfinite(temps[0]):
        ax2.set_yticklabels([f"{temps[0]:.1f}" for _ in yticks])
    else:
        ax2.set_yticklabels([""] * len(yticks))

    title = (
        f"CA Grid: Lcomp, X=[{xmin},{xmax}], Y=[{ymin},{ymax}], "
        f"stride=({xstride},{ystride}), mode={mode}"
    )
    # ax.set_title(title)

    do_save, save_dpi, out_name = ask_save_dpi(preset["save_dpi"])
    if do_save:
        if out_name is None:
            out_name = build_output_name(path, mode, tile2x)
        fig.savefig(out_name, dpi=save_dpi, bbox_inches="tight")
        print(f"Saved figure to: {out_name}  [dpi={save_dpi}]")

    plt.show()


if __name__ == "__main__":
    main()