"""tkplot.py

フィット結果の標準プロット補助。

細かい見た目はアプリ側で調整できるよう、
Figure/Axes を返す薄い関数にしている。
"""

from __future__ import annotations

from pathlib import Path
from typing import Callable, Mapping, Optional, Sequence, Tuple, Union
import numpy as np


ArrayLike = Sequence[float] | np.ndarray


def make_xcal(
    x: ArrayLike,
    *,
    n: int = 401,
    xmin: Optional[float] = None,
    xmax: Optional[float] = None,
    margin: float = 0.0,
) -> np.ndarray:
    """プロット用の滑らかな x 軸を作る。"""
    x_arr = np.asarray(x, dtype=float).reshape(-1)

    if xmin is None:
        xmin = float(np.nanmin(x_arr))
    if xmax is None:
        xmax = float(np.nanmax(x_arr))

    width = xmax - xmin
    xmin = xmin - margin * width
    xmax = xmax + margin * width

    return np.linspace(xmin, xmax, int(n))


def plot_fit_before_after(
    x: ArrayLike,
    y: ArrayLike,
    model_func: Callable[[np.ndarray, Mapping[str, float]], ArrayLike],
    p_before: Mapping[str, float],
    *,
    p_after: Optional[Mapping[str, float]] = None,
    xcal: Optional[ArrayLike] = None,
    yerr: Optional[ArrayLike] = None,
    band: Optional[Mapping[str, ArrayLike]] = None,
    xlabel: str = "x",
    ylabel: str = "y",
    title: str = "fit result",
    data_label: str = "data",
    before_label: str = "before",
    after_label: str = "after",
    out_png: Optional[Union[str, Path]] = None,
    show: bool = False,
    close: bool = True,
):
    """データ、フィット前、フィット後を重ねて描画する。

    model_func:
        y = model_func(x_array, params_dict)

    band:
        {"x": xband, "y_low": y_low, "y_high": y_high} または
        {"x": xband, "y_mean": y_mean, "sigma": sigma}
    """
    import matplotlib.pyplot as plt

    x_arr = np.asarray(x, dtype=float).reshape(-1)
    y_arr = np.asarray(y, dtype=float).reshape(-1)

    if xcal is None:
        xcal_arr = make_xcal(x_arr)
    else:
        xcal_arr = np.asarray(xcal, dtype=float).reshape(-1)

    fig, ax = plt.subplots(figsize=(8, 5))

    if yerr is None:
        ax.scatter(x_arr, y_arr, s=28, color="black", alpha=0.75, label=data_label)
    else:
        ax.errorbar(
            x_arr,
            y_arr,
            yerr=np.asarray(yerr, dtype=float).reshape(-1),
            fmt="o",
            ms=4,
            color="black",
            alpha=0.75,
            label=data_label,
        )

    y_before = np.asarray(model_func(xcal_arr, p_before), dtype=float).reshape(-1)
    ax.plot(xcal_arr, y_before, "--", color="tab:orange", lw=2, label=before_label)

    if p_after is not None:
        y_after = np.asarray(model_func(xcal_arr, p_after), dtype=float).reshape(-1)
        ax.plot(xcal_arr, y_after, "-", color="tab:blue", lw=2, label=after_label)

    if band is not None:
        plot_band(ax, band)

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    ax.legend()
    fig.tight_layout()

    if out_png is not None:
        out_png = Path(out_png)
        out_png.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(out_png, dpi=160)

    if show:
        plt.show()

    if close:
        plt.close(fig)

    return fig, ax


def plot_band(
    ax,
    band: Mapping[str, ArrayLike],
    *,
    color: str = "tab:blue",
    alpha: float = 0.18,
    label: str = "uncertainty",
):
    """既存Axesに誤差帯を追加する。"""
    if "x" not in band:
        raise ValueError("band must contain 'x'")

    x = np.asarray(band["x"], dtype=float).reshape(-1)

    if "y_low" in band and "y_high" in band:
        y_low = np.asarray(band["y_low"], dtype=float).reshape(-1)
        y_high = np.asarray(band["y_high"], dtype=float).reshape(-1)
    elif "y_mean" in band and "sigma" in band:
        y_mean = np.asarray(band["y_mean"], dtype=float).reshape(-1)
        sigma = np.asarray(band["sigma"], dtype=float).reshape(-1)
        y_low = y_mean - sigma
        y_high = y_mean + sigma
    else:
        raise ValueError("band must contain either y_low/y_high or y_mean/sigma")

    ax.fill_between(x, y_low, y_high, color=color, alpha=alpha, label=label)


def save_progress_plot(
    iteration: int,
    x: ArrayLike,
    y: ArrayLike,
    model_func: Callable[[np.ndarray, Mapping[str, float]], ArrayLike],
    p_before: Mapping[str, float],
    p_current: Mapping[str, float],
    *,
    out_dir: Union[str, Path] = ".",
    prefix: str = "fit_progress",
    **kwargs,
) -> Path:
    """フィット途中の画像を保存する。callback から呼ぶ想定。"""
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    out_png = out_dir / f"{prefix}_{int(iteration):04d}.png"

    plot_fit_before_after(
        x,
        y,
        model_func,
        p_before,
        p_after=p_current,
        out_png=out_png,
        title=f"fit progress iter={iteration}",
        **kwargs,
    )

    return out_png