"""
=============================================================================
CS230 · 第 3 讲工具模块：把二维数据和决策边界画出来（终端字符画 + 可选的 PNG）
=============================================================================

对应讲义：lectures/lecture03_shallow_neural_networks.md（§1、§2.4、§8 的"字符画"都由这里生成）

-----------------------------------------------------------------------------
【字符画怎么读】
-----------------------------------------------------------------------------
把平面切成 width × height 个小格子（默认 40 × 20；终端里一个字符大约"高是宽的两倍"，
所以 40 × 20 看起来接近正方形）。每个格子：

    背景  ' '（空白）  模型在格子中心预测为第 0 类
          '.'          模型在格子中心预测为第 1 类
    数据点 'o'          落在这个格子里的样本（多数）是第 0 类，而且格子被预测为第 0 类（分对了）
          'x'          落在这个格子里的样本（多数）是第 1 类，而且格子被预测为第 1 类（分对了）
          'O' / 'X'    **大写 = 分错了**：第 0 类的点落在 '.' 区域，或第 1 类的点落在空白区域
          '*'          这个格子里两类的点一样多

所以一眼看过去：空白和 '.' 的分界线就是决策边界；大写字母越少，分得越好。

【PNG】装了 matplotlib 时，save_boundary_grid / save_activation_figure 会把同样的内容画成图片
（讲义里引用的 figures/*.png）；没装时返回 False，不报错。

【注释标签】【思路】【参数】【返回】【边界】【numpy 小知识】【Python 小知识】
"""

import os
import sys

import numpy as np

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import pad  # noqa: E402  （中英文混排时按显示宽度补空格；第 1 讲的小工具）


# =============================================================================
# 1. 终端字符画
# =============================================================================
def data_bounds(X, margin=0.5):
    """
    画图的范围：数据在两个坐标上的最小值 / 最大值，再往外留 margin。

    【参数】X：(2, m)；margin：留白
    【返回】(x1_min, x1_max, x2_min, x2_max)，四个 float
    """
    return (float(X[0].min() - margin), float(X[0].max() + margin),
            float(X[1].min() - margin), float(X[1].max() + margin))


def grid_centers(bounds, width, height):
    """
    每个格子中心的坐标，按"从上到下、从左到右"的顺序排成 (2, width·height) 的矩阵。

    【思路】第 0 行字符在最上面，对应 x₂ 最大；所以 x₂ 从大到小排（和数学里"y 轴朝上"一致）。
    【参数】bounds：data_bounds 的返回值；width、height：格子数
    【返回】G：(2, height·width)。第 r 行第 c 列的格子是 G 的第 r·width + c 列。
    【numpy 小知识】np.meshgrid(a, b) 生成两个二维数组：第一个每一行都是 a，第二个每一列都是 b。
            .ravel() 把多维数组按行拉平成一维。
    """
    x1_min, x1_max, x2_min, x2_max = bounds
    dx = (x1_max - x1_min) / width
    dy = (x2_max - x2_min) / height
    xs = x1_min + dx * (np.arange(width) + 0.5)                 # 格子中心的 x₁
    ys = x2_max - dy * (np.arange(height) + 0.5)                # 格子中心的 x₂（从上往下变小）
    gx, gy = np.meshgrid(xs, ys)
    return np.vstack([gx.ravel(), gy.ravel()])


def ascii_boundary(predict_fn, X=None, Y=None, bounds=None, width=40, height=20, show_points=True):
    """
    画一张字符画：背景是模型的预测区域，上面叠上数据点（分错的点用大写字母）。

    【思路】1. 对每个格子的中心点调用 predict_fn，得到 0 / 1 → 背景字符 ' ' / '.'。
            2. 把每个数据点放进它所在的格子；一个格子里有多个点时按多数类决定字符。
            3. 加一个边框。
    【参数】
        predict_fn  : 函数，输入 (2, k) 的坐标，返回 (1, k) 或 (k,) 的 0 / 1 预测
        X, Y        : 要叠加的数据点 (2, m)、(1, m)；show_points=False 或 X 为 None 时只画背景
        bounds      : 画图范围；None 时用 data_bounds(X)
        width, height : 字符画的列数、行数
    【返回】list[str]：每个元素是一行（含边框），可以直接 print，也可以交给 side_by_side 并排打印
    【边界】数据点落在 bounds 之外时被忽略（用 data_bounds 时不会发生）。
    【Python 小知识】函数也可以当参数传进来（predict_fn）。调用方常用 lambda 临时写一个：
            lambda G: predict(params, G)  等价于  def f(G): return predict(params, G)
    """
    if bounds is None:
        bounds = data_bounds(X)
    G = grid_centers(bounds, width, height)
    pred = np.asarray(predict_fn(G)).ravel().reshape(height, width)       # (height, width) 的 0/1
    canvas = [["." if pred[r, c] > 0.5 else " " for c in range(width)] for r in range(height)]

    if show_points and X is not None:
        x1_min, x1_max, x2_min, x2_max = bounds
        cols = np.floor((X[0] - x1_min) / (x1_max - x1_min) * width).astype(int)
        rows = np.floor((x2_max - X[1]) / (x2_max - x2_min) * height).astype(int)
        count = np.zeros((height, width, 2), dtype=int)                  # 每个格子里两类各有几个点
        labels = Y.ravel().astype(int)
        for r, c, y in zip(rows, cols, labels):
            if 0 <= r < height and 0 <= c < width:
                count[r, c, y] += 1
        for r in range(height):
            for c in range(width):
                n0, n1 = count[r, c]
                if n0 + n1 == 0:
                    continue
                if n0 == n1:
                    canvas[r][c] = "*"
                    continue
                cls = 1 if n1 > n0 else 0
                wrong = (pred[r, c] > 0.5) != (cls == 1)                 # 点的类别和格子的预测不一致
                ch = "x" if cls == 1 else "o"
                canvas[r][c] = ch.upper() if wrong else ch
    border = "+" + "-" * width + "+"
    return [border] + ["|" + "".join(row) + "|" for row in canvas] + [border]


def side_by_side(panels, titles, gap=3):
    """
    把几张字符画左右并排拼成若干行文字，每张上面一行标题。

    【参数】panels：list，每个元素是 ascii_boundary 的返回值（行数可以不同）；
            titles：和 panels 一样长的标题列表（可以含中文）；gap：两张之间的空格数
    【返回】list[str]
    【Python 小知识】max(len(p) for p in panels) 里的 (... for ...) 是**生成器表达式**：
            像列表推导式一样逐个产生值，但不先建一个列表。
    """
    widths = [max(len(line) for line in p) for p in panels]
    height = max(len(p) for p in panels)
    out = [(" " * gap).join(pad(t, w) for t, w in zip(titles, widths)).rstrip()]
    for i in range(height):
        parts = [(p[i] if i < len(p) else "").ljust(w) for p, w in zip(panels, widths)]
        out.append((" " * gap).join(parts).rstrip())
    return out


def print_lines(lines):
    """逐行打印一个字符串列表。"""
    for line in lines:
        print(line)


ASCII_LEGEND = ("图例：背景 ' ' = 预测第 0 类，'.' = 预测第 1 类；点 o = 第 0 类，x = 第 1 类；"
                "大写 O / X = 分错的点；* = 两类的点挤在同一格")


# =============================================================================
# 2. 可选的 PNG（需要 matplotlib）
# =============================================================================
# 颜色：两类数据用同一组"类别色"——第 0 类蓝、第 1 类橙（和第 1 讲的图同一套颜色）；
# 预测区域用这两种颜色的浅色调，数据点用实色。
COLOR_CLASS = ("#2a78d6", "#eb6834")
COLOR_REGION = ("#d9e8fb", "#fbe1d5")
COLOR_SURFACE = "#fcfcfb"
COLOR_TEXT = "#0b0b0b"
COLOR_TEXT_2 = "#52514e"
COLOR_GRID = "#e6e5e0"
COLOR_AXIS = "#c3c2b7"


def _setup_matplotlib():
    """
    导入 matplotlib 并设置中文字体；没装时返回 None。

    【Python 小知识】try / except ImportError：导入失败时执行 except 里的代码，而不是让程序崩溃。
    """
    try:
        import matplotlib
        matplotlib.use("Agg")                        # 不弹窗口，只画到文件里
        import matplotlib.pyplot as plt
    except ImportError:
        return None
    plt.rcParams["font.sans-serif"] = ["PingFang SC", "Heiti SC", "Arial Unicode MS", "Noto Sans CJK SC",
                                       "Microsoft YaHei", "DejaVu Sans"]
    plt.rcParams["axes.unicode_minus"] = False
    return plt


def save_boundary_grid(panels, X, Y, path, ncols=4, resolution=200, suptitle=None):
    """
    把几个模型的决策边界画成一张多宫格 PNG（Coursera 作业里那张"隐藏单元数 1–50"的图）。

    【参数】
        panels : list of (title, predict_fn)
        X, Y   : 叠加的数据点（通常是训练集）
        path   : 保存路径
        ncols  : 每行几张小图
        resolution : 每张小图背景网格的分辨率（resolution × resolution 个点）
        suptitle   : 整张图的标题
    【返回】True 表示保存成功；没装 matplotlib 时返回 False
    """
    plt = _setup_matplotlib()
    if plt is None:
        return False
    from matplotlib.colors import ListedColormap
    bounds = data_bounds(X)
    G = grid_centers(bounds, resolution, resolution)
    nrows = int(np.ceil(len(panels) / ncols))
    fig, axes = plt.subplots(nrows, ncols, figsize=(3.1 * ncols, 3.2 * nrows), dpi=130, squeeze=False)
    fig.patch.set_facecolor(COLOR_SURFACE)
    labels = Y.ravel()
    for k, ax in enumerate(axes.ravel()):
        if k >= len(panels):
            ax.axis("off")
            continue
        title, fn = panels[k]
        pred = np.asarray(fn(G)).ravel().reshape(resolution, resolution)
        ax.imshow(pred, extent=bounds, origin="upper", cmap=ListedColormap(COLOR_REGION),
                  vmin=0, vmax=1, interpolation="nearest", aspect="auto")
        for cls in (0, 1):
            sel = labels == cls
            ax.scatter(X[0, sel], X[1, sel], s=9, color=COLOR_CLASS[cls], edgecolors=COLOR_SURFACE,
                       linewidths=0.4, label=f"第 {cls} 类")
        ax.set_title(title, fontsize=10, color=COLOR_TEXT)
        ax.set_xticks([])
        ax.set_yticks([])
        for side in ax.spines.values():
            side.set_color(COLOR_AXIS)
    handles, names = axes.ravel()[0].get_legend_handles_labels()
    fig.legend(handles, names, loc="lower center", ncol=2, frameon=False, fontsize=9,
               labelcolor=COLOR_TEXT_2, bbox_to_anchor=(0.5, 0.0))
    if suptitle:
        fig.suptitle(suptitle, color=COLOR_TEXT, fontsize=11)
    fig.tight_layout(rect=(0, 0.04, 1, 1))
    os.makedirs(os.path.dirname(path), exist_ok=True)
    fig.savefig(path, facecolor=fig.get_facecolor())
    plt.close(fig)
    return True


def save_activation_figure(path, funcs):
    """
    把几个激活函数和它们的导数画成一排小图（每张图：实线 = g(z)，虚线 = g'(z)）。

    【参数】path：保存路径；funcs：list of (名字, g, g')，g 和 g' 都是"输入 z 数组、返回同形状数组"的函数
    【返回】True / False（没装 matplotlib）
    """
    plt = _setup_matplotlib()
    if plt is None:
        return False
    z = np.linspace(-5, 5, 1001)
    fig, axes = plt.subplots(1, len(funcs), figsize=(3.0 * len(funcs), 3.0), dpi=130, sharey=True)
    fig.patch.set_facecolor(COLOR_SURFACE)
    for ax, (name, g, dg) in zip(axes, funcs):
        ax.set_facecolor(COLOR_SURFACE)
        ax.axhline(0, color=COLOR_AXIS, linewidth=0.8)
        ax.axvline(0, color=COLOR_AXIS, linewidth=0.8)
        ax.plot(z, g(z), color=COLOR_CLASS[0], linewidth=2, label="g(z)")
        ax.plot(z, dg(z), color=COLOR_CLASS[1], linewidth=2, linestyle="--", label="g'(z)")
        ax.set_title(name, fontsize=11, color=COLOR_TEXT)
        ax.set_ylim(-1.3, 2.3)
        ax.set_xlabel("z", color=COLOR_TEXT_2)
        ax.grid(True, color=COLOR_GRID, linewidth=0.8)
        for side in ("top", "right"):
            ax.spines[side].set_visible(False)
        for side in ("left", "bottom"):
            ax.spines[side].set_color(COLOR_AXIS)
        ax.tick_params(colors=COLOR_TEXT_2, labelsize=8)
    axes[0].legend(loc="upper left", frameon=False, fontsize=9, labelcolor=COLOR_TEXT_2)
    fig.tight_layout()
    os.makedirs(os.path.dirname(path), exist_ok=True)
    fig.savefig(path, facecolor=fig.get_facecolor())
    plt.close(fig)
    return True