graph_renderer.py

# graph_renderer.py
import io
import base64

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import yaml
from typing import Optional, Dict, Any


def parse_frontmatter(md_text: str) -> Dict[str, Any]:
    """
    Extract frontmatter dict from md_text. Supports YAML delimited by ---
    or a top contiguous key: value block (until first blank line).
    """
    md = md_text.lstrip()
    if md.startswith('---'):
        # YAML frontmatter block
        end = md.find('\n---', 3)
        if end != -1:
            block = md[3:end]
        else:
            # no closing marker; consider until first blank line
            lines = md.splitlines()
            for i, line in enumerate(lines[1:], start=1):
                if line.strip() == '':
                    block = '\n'.join(lines[1:i])
                    break
            else:
                block = '\n'.join(lines[1:])
    else:
        # contiguous key:value lines at top
        lines = md.splitlines()
        block_lines = []
        for line in lines:
            if line.strip() == '':
                break
            if ':' in line:
                block_lines.append(line)
            else:
                break
        block = '\n'.join(block_lines)

    if not block.strip():
        return {}
    try:
        data = yaml.safe_load(block)
        return data or {}
    except Exception:
        return {}


def _validate_params(params: Dict[str, Any]) -> Dict[str, Any]:
    """
    Ensure and coerce expected params: n (int), omega (float), optional coeffs (list of floats).
    Enforce 1 <= n <= 10.
    Returns normalized dict or raises ValueError.
    """
    if not isinstance(params, dict):
        raise ValueError("Invalid params")

    n = params.get('n', None)
    if n is None:
        raise ValueError("Missing 'n' parameter for graph")
    try:
        n = int(n)
    except Exception:
        raise ValueError("'n' must be an integer")
    if n < 1 or n > 10:
        raise ValueError("'n' must be between 1 and 10")

    omega = params.get('omega', 1.0)
    try:
        omega = float(omega)
    except Exception:
        omega = 1.0

    coeffs = params.get('coeffs', None)
    if coeffs is not None:
        # try to coerce to list of floats
        try:
            coeffs = [float(c) for c in coeffs]
        except Exception:
            coeffs = None

    wave = params.get('wave', 'triangle')
    graph_type = params.get('graph_type')
    if graph_type == 'saw_wave':
        wave = 'saw'

    return {'n': n, 'omega': omega, 'coeffs': coeffs, 'wave': wave}


def _compute_exponential_fourier(n: int, omega: float, coeffs: Optional[list], t: np.ndarray, wave: str = "triangle") -> np.ndarray:
    """
    Compute exponential-form partial sum:
    S(t) = sum_{k=-n, k!=0}^{n} a_k * exp(j * k * omega * t) + optional DC
    Supports 'triangle' and 'square' wave coefficients.
    """
    ks = np.arange(-n, n + 1)
    ks = ks[ks != 0]  # exclude k=0

    if coeffs is not None:
        # assume coeffs correspond to ks in order; coerce to complex
        a_k = np.array([complex(c) for c in coeffs], dtype=complex)
        if a_k.shape[0] != ks.shape[0]:
            # ignore provided coeffs if length mismatch
            a_k = None
        else:
            a_k = a_k
    else:
        a_k = None

    if a_k is None:
        freq_multiplier = omega
        if wave == "triangle":
            a_k = (np.power(-1.0, ks) - 1.0) / (ks.astype(float) ** 2 * (np.pi ** 2))
            a_k = a_k.astype(float)  # real coefficients
        elif wave == "square":
            # a_k = 2/(j*pi*k) for odd k, 0 for even k
            a_k = np.zeros_like(ks, dtype=complex)
            odd_mask = (np.abs(ks) % 2) == 1
            a_k[odd_mask] = 2.0 / (1j * np.pi * ks[odd_mask].astype(float))
        elif wave == "saw":
            # coefficients from exponential-form saw wave
            a_k = 1j / (np.pi * ks.astype(float))
            freq_multiplier = omega * np.pi
        else:
            # default fallback to triangle
            a_k = (np.power(-1.0, ks) - 1.0) / (ks.astype(float) ** 2 * (np.pi ** 2))
            a_k = a_k.astype(float)

    # vectorized complex exponential sum
    exponents = np.exp(1j * np.outer(ks, freq_multiplier * t))
    s = a_k[:, None] * exponents
    total = s.sum(axis=0)
    # For triangle we have a DC 1/2 term in previous impl; keep that if wave is triangle
    if wave == "triangle":
        result = 0.5 + np.real(total)
    elif wave == "saw":
        result = 1.0 + np.real(total)
    else:
        result = np.real(total)
    return result


def render_fourier_graph_png_bytes(params: Dict[str, Any], width: int = 800, height: int = 320) -> bytes:
    """
    Validate params, compute series on t grid, render PNG bytes and return them.
    """
    p = _validate_params(params)
    n = p['n']
    omega = p['omega']
    coeffs = p.get('coeffs', None)
    wave = p.get('wave', 'triangle')

    # t grid: one period [0, 2*pi)
    t = np.linspace(0, 2 * np.pi, 800)
    y = _compute_exponential_fourier(n=n, omega=omega, coeffs=coeffs, t=t, wave=wave)

    fig = plt.figure(figsize=(width / 100.0, height / 100.0), dpi=100)
    ax = fig.add_subplot(111)
    ax.plot(t, y, lw=1.5)
    ax.set_xlabel('t')
    ax.set_ylabel('f(t)')
    ax.grid(True, alpha=0.3)
    ax.set_title(f'Fourier partial sum ({wave}, n={n}, ω={omega})')
    plt.tight_layout()

    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100)
    plt.close(fig)
    buf.seek(0)
    return buf.read()


def render_fourier_data_url(params: Dict[str, Any]) -> str:
    """
    Return a data URL string 'data:image/png;base64,...' or an HTML error snippet on failure.
    """
    try:
        png = render_fourier_graph_png_bytes(params)
        b64 = base64.b64encode(png).decode('ascii')
        return 'data:image/png;base64,' + b64
    except Exception as e:
        # return an HTML-safe error block
        msg = f"Graph error: {str(e)}"
        return f'<div class="graph-error">Graph error: {msg}</div>'