editor.py

"""
Matplotlib Theme Editor Streamlit Application
============================================

This Streamlit app provides an interactive interface for configuring
Matplotlib's ``rcParams``. Users can browse and edit all parameters,
search for specific settings, import/export themes in JSON or YAML
formats, reset to defaults, choose a global colour map, register custom
fonts, and preview the effects on five representative plots
(line, scatter, bar, pie, radar).

Run:
    pip install -r requirements.txt
    streamlit run app.py
"""
from __future__ import annotations

import ast
import io
import json
from typing import Any, Dict, Iterable, List, Tuple

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import streamlit as st
from matplotlib.colors import is_color_like

from theme_manager import (
    apply_overrides,
    enum_to_str,
    export_theme_dict,
    get_all_rcparams,
    group_rcparams,
    list_cmaps,
    load_theme,
    normalize_value_for_rc,
    register_fonts,
    reset_to_defaults,
    save_theme,
)

# ──────────────────────────────────────────────────────────────────────────────
# Page setup
# ──────────────────────────────────────────────────────────────────────────────
st.set_page_config(page_title="Matplotlib Theme Editor", layout="wide")
st.title("🎨 Matplotlib Theme Editor")
st.caption(
    "Edit rcParams • Import/Export JSON/YAML • Live previews • Upload custom fonts"
)

# Sidebar: global actions
st.sidebar.header("Theme Controls")
col1, col2 = st.sidebar.columns(2)
with col1:
    if st.button("Reset to defaults", use_container_width=True):
        reset_to_defaults()
        st.session_state.pop("overrides", None)
        st.success("Reset Matplotlib defaults.")
        st.rerun()
with col2:
    if st.button("Clear session overrides", use_container_width=True):
        st.session_state.pop("overrides", None)
        st.info("Session overrides cleared.")
        st.rerun()

# Import/export
with st.sidebar.expander("⬆️ Import Theme (JSON or YAML)", expanded=False):
    up = st.file_uploader("Choose a theme file", type=["json", "yaml", "yml", "mplstyle"])
    if up is not None:
        try:
            overrides = load_theme(up)
            st.session_state.setdefault("overrides", {}).update(overrides)
            st.success(f"Loaded {len(overrides)} overrides.")
            # apply immediately
            applied, errors = apply_overrides(overrides)
            if applied:
                st.success(f"Applied {len(applied)} parameter(s).")
            if errors:
                with st.expander("⚠️ Some parameters failed", expanded=False):
                    for e in errors:
                        st.write("- ", e)
        except Exception as e:
            st.error(f"Import failed: {e}")

with st.sidebar.expander("⬇️ Export Theme", expanded=False):
    export_fmt = st.selectbox("Format", ["json", "yaml", "mplstyle"])
    export_name = st.text_input("Filename", value=f"theme.{export_fmt}")
    if st.button("Download", use_container_width=True):
        data = export_theme_dict(mpl.rcParams)
        if export_fmt == "json":
            byts = io.BytesIO(json.dumps(data, indent=2).encode("utf-8"))
            st.download_button("Download JSON", byts, file_name=export_name, mime="application/json")
        elif export_fmt in ("yaml", "yml"):
            import yaml
            byts = io.BytesIO(yaml.safe_dump(data, sort_keys=True, allow_unicode=True).encode("utf-8"))
            st.download_button("Download YAML", byts, file_name=export_name, mime="text/yaml")
        else:  # mplstyle
            lines = []
            for k, v in data.items():
                if isinstance(v, list):
                    v = tuple(v)
                lines.append(f"{k}: {v}")
            byts = io.BytesIO("\n".join(lines).encode("utf-8"))
            st.download_button("Download .mplstyle", byts, file_name=export_name, mime="text/plain")

# Font upload
with st.sidebar.expander("🔤 Upload & Register Custom Fonts (.ttf/.otf)", expanded=False):
    uploaded_fonts = st.file_uploader(
        "Upload font files", type=["ttf", "otf"], accept_multiple_files=True
    )
    if uploaded_fonts:
        try:
            names = register_fonts(uploaded_fonts)
            if names:
                st.success("Registered: " + ", ".join(names))
            else:
                st.info("No usable fonts found.")
        except Exception as e:
            st.error(f"Font registration failed: {e}")

# Colormap control (applies to image.cmap)
with st.sidebar.expander("🎨 Default Colormap", expanded=True):
    current = mpl.rcParams.get("image.cmap", "viridis")
    cmap = st.selectbox("image.cmap", list_cmaps(), index=list_cmaps().index(str(current)) if str(current) in list_cmaps() else 0)
    if cmap != current:
        mpl.rcParams["image.cmap"] = cmap
        st.session_state.setdefault("overrides", {})["image.cmap"] = cmap

# Search/filter controls
with st.sidebar.expander("🔎 Search & Filter", expanded=True):
    query = st.text_input("Filter by substring", "")
    show_all = st.checkbox("Show raw editor for all", value=False)

# ──────────────────────────────────────────────────────────────────────────────
# Data & helpers
# ──────────────────────────────────────────────────────────────────────────────
all_params = get_all_rcparams()
if "overrides" in st.session_state:
    # show pending session overrides visually (without applying yet)
    for k, v in st.session_state["overrides"].items():
        all_params[k] = v

# group for ergonomic browsing
ordered_groups = (
    [
        "figure",
        "axes",
        "lines",
        "font",
        "text",
        "legend",
        "grid",
        "xtick",
        "ytick",
        "patch",
        "image",
        "savefig",
        "scatter",
        "hist",
    ]
)

by_group = group_rcparams(all_params)

# Widget factory
from matplotlib import rcParamsDefault
import enum as _enum

all_keys_id = []

def sanitize_color(c: str) -> str:
    if c=="white":
        return "#ffffff"
    if c=="black":
        return "#000000"
    if c=="C0":
        return "#1f77b4"
    if c=="C1":
        return "#ff7f0e"
    if c=="C2":
        return "#2ca02c"
    if c=="C3":
        return "#d62728"
    if c=="C4":
        return "#9467bd"
    if c=="C5":
        return "#8c564b"
    return c

def check_correct_color(c: str) -> bool:
    if not isinstance(c, str):
        return False
    if c in ["none", "None", "AUTO", "auto"]:
        return False
    if not c.startswith("#"):
        return False
    try:
        return is_color_like(c)
    except Exception:
        return False

def make_widget(key: str, cur: Any):
    """Return a Streamlit widget for a single rcParam and the chosen value.
    Handles colours, bools, numbers, enums, known string choices, and
    (list, tuple) via literal syntax.
    """
    cur = sanitize_color(cur)
    if key in all_keys_id:
        key = f"{key}_{all_keys_id.count(key)}"
    all_keys_id.append(key)

    # colour picker
    if (key.endswith("color") or key in {
        "axes.facecolor","axes.edgecolor","axes.labelcolor","figure.facecolor",
        "figure.edgecolor","text.color","xtick.color","ytick.color","grid.color",
        "patch.facecolor","patch.edgecolor","savefig.facecolor","savefig.edgecolor",
        "legend.edgecolor","legend.facecolor",
    }) and isinstance(cur, str) and cur != "auto":
        return st.color_picker(key, value=cur if check_correct_color(cur) else "#000000")

    # bool
    if isinstance(cur, bool):
        return st.checkbox(key, value=cur)

    # enum (based on default type)
    default = rcParamsDefault.get(key, None)
    if isinstance(default, _enum.Enum):
        choices = [e.name.lower() for e in type(default)]
        cur_name = enum_to_str(cur)
        idx = choices.index(cur_name) if cur_name in choices else 0
        return st.selectbox(key, choices, index=idx)

    # known choice lists
    CHOICES = {
        "axes.titlelocation": ["center", "left", "right"],
        "toolbar": ["toolbar2", "toolmanager", "None"],
        "mathtext.fontset": ["dejavusans","dejavuserif","cm","stix","stixsans","custom"],
        "image.interpolation": [
            "none","nearest","antialiased","bilinear","bicubic","spline16","spline36",
            "hanning","hamming","hermite","kaiser","quadric","catrom","gaussian",
            "bessel","mitchell","sinc","lanczos",
        ],
    }
    if key in CHOICES:
        choices = CHOICES[key]
        cur_str = str(cur)
        idx = choices.index(cur_str) if cur_str in choices else 0
        return st.selectbox(key, choices, index=idx)

    # number
    if isinstance(cur, (int, float)):
        return st.number_input(key, value=float(cur))

    # string
    if isinstance(cur, str):
        return st.text_input(key, value=cur)

    # collection
    if isinstance(cur, (list, tuple)):
        return st.text_input(key, value=repr(cur))

    return st.text_input(key, value=str(cur))

# ──────────────────────────────────────────────────────────────────────────────
# UI layout
# ──────────────────────────────────────────────────────────────────────────────
updated: Dict[str, Any] = {}

side_tab = st.sidebar.radio(
    "Choose editor view:",
    ["Ergonomic", "Search"],
    label_visibility="collapsed"
)

updated: Dict[str, Any] = {}

if side_tab == "Ergonomic":
    st.sidebar.subheader("Ergonomic groups")
    for grp in ordered_groups + [g for g in sorted(by_group) if g not in ordered_groups]:
        if grp not in by_group:
            continue
        with st.sidebar.expander(grp):
            for key, cur in sorted(by_group[grp], key=lambda kv: kv[0]):
                if query and query.lower() not in key.lower():
                    continue
                updated[key] = make_widget(key, cur)

elif side_tab == "Search":
    st.sidebar.subheader("Search all rcParams")
    cnt = 0
    for key in sorted(all_params):
        if query and query.lower() not in key.lower():
            continue
        updated[key] = make_widget(key, all_params[key])
        cnt += 1
    if query:
        st.sidebar.info(f"Showing {cnt} parameter(s) matching '{query}'.")

# Apply controls
apply_col1, apply_col2 = st.columns([1,1])
with apply_col1:
    do_apply = st.button("Apply changes", type="primary")
with apply_col2:
    remember = st.checkbox("Remember in session", value=True)

if do_apply:
    # coerce string literals and normalize for rc
    prepared: Dict[str, Any] = {}
    for k, v in updated.items():
        if isinstance(v, str) and v.strip().startswith(("(", "[", "{", "'", '"')):
            try:
                v = ast.literal_eval(v)
            except Exception:
                pass
        v = normalize_value_for_rc(k, v)
        prepared[k] = v
    applied, errors = apply_overrides(prepared)
    if remember:
        st.session_state.setdefault("overrides", {}).update(applied)
    if applied:
        st.success(f"Applied {len(applied)} parameter(s).")
    if errors:
        with st.expander("⚠️ Some parameters failed", expanded=False):
            for e in errors:
                st.write("- ", e)

# demo data
x = np.linspace(0, np.pi, 200)
y = np.sin(2*x)
rng = np.random.default_rng(0)

def fig_line():
    f = plt.figure()
    ax = f.add_subplot(1, 1, 1)
    ax.plot(x, y, label="sin(2x)")
    ax.plot(x, np.sin(3*x), label="sin(3x)")
    ax.plot(x, np.sin(4 * x), label="sin(4x)")
    ax.plot(x, np.sin(5 * x), label="sin(5x)")
    ax.plot(x, np.sin(6 * x), label="sin(6x)")
    ax.set_title("Line Plot")
    ax.legend(); ax.grid(True)
    # set labels
    ax.set_xlabel("X axis"); ax.set_ylabel("Y axis")
    return f

def fig_scatter():
    f = plt.figure()
    ax = f.add_subplot(1, 1, 1)
    xs = rng.normal(0,1,200); ys = rng.normal(0,1,200)
    ax.scatter(xs, ys, c=ys, cmap=mpl.rcParams.get("image.cmap", "viridis"))
    ax.set_title("Scatter Plot"); ax.grid(True)
    # set labels
    ax.set_xlabel("X axis"); ax.set_ylabel("Y axis")
    return f

def fig_bar():
    f = plt.figure()
    ax = f.add_subplot(1, 1, 1)
    cats = ["A","B","C","D","E"]; vals = rng.uniform(0.5,1.5,len(cats))
    ax.bar(cats, vals); ax.set_title("Bar Plot"); ax.grid(True, axis="y")
    # set labels
    ax.set_xlabel("Category"); ax.set_ylabel("Value")
    return f

def fig_pie():
    f = plt.figure()
    ax = f.add_subplot(1, 1, 1)
    vals = np.array([15,30,25,20,10]); labels=["A","B","C","D","E"]
    ax.pie(vals, labels=labels, autopct="%1.0f%%", startangle=90)
    ax.set_title("Pie Chart"); ax.axis("equal"); return f

def fig_radar():
    cats=["A","B","C","D","E"]; vals=rng.uniform(0.2,1.0,len(cats))
    ang = np.linspace(0, 2*np.pi, len(cats), endpoint=False)
    vals = np.concatenate((vals,[vals[0]])); ang = np.concatenate((ang,[ang[0]]))
    f = plt.figure()
    f.set_figheight(4)
    ax = f.add_subplot(1, 1, 1, polar=True)
    ax.plot(ang, vals, "o-", linewidth=2); ax.fill(ang, vals, alpha=0.25)
    ax.set_thetagrids(np.degrees(ang[:-1]), cats); ax.set_title("Radar Plot")
    return f

def fig_boxplot():
    f = plt.figure()
    ax = f.add_subplot(1, 1, 1)
    rng = np.random.default_rng(0)
    data = [rng.normal(loc, 0.3*loc, 300) for loc in range(5)]
    ax.boxplot(data, notch=True, patch_artist=True)
    ax.set_title("Box Plot"); ax.set_xticklabels(["A","B","C","D","E"])
    # set labels
    ax.set_xlabel("Category"); ax.set_ylabel("Value")
    return f

# Preview & export
st.subheader("Live preview")

colL, colC, colR = st.columns(3)

with colL:
    st.pyplot(fig_line()); st.pyplot(fig_bar())
with colC:
    st.pyplot(fig_boxplot());    st.pyplot(fig_radar())
with colR:
    st.pyplot(fig_scatter()); st.pyplot(fig_pie())

st.markdown("---")
st.write("Quick export buttons (JSON/YAML/.mplstyle are also in sidebar):")
data = export_theme_dict(mpl.rcParams)
json_bytes = io.BytesIO(json.dumps(data, indent=2).encode("utf-8"))
st.download_button("Download JSON", json_bytes, file_name="theme.json", mime="application/json")

import yaml
yaml_bytes = io.BytesIO(yaml.safe_dump(data, sort_keys=True, allow_unicode=True).encode("utf-8"))
st.download_button("Download YAML", yaml_bytes, file_name="theme.yaml", mime="text/yaml")

lines = []
for k, v in data.items():
    if isinstance(v, list):
        v = tuple(v)
    lines.append(f"{k}: {v}")
style_bytes = io.BytesIO("\n".join(lines).encode("utf-8"))
st.download_button("Download .mplstyle", style_bytes, file_name="theme.mplstyle", mime="text/plain")