import os

import numpy as np
import torch
from torch_geometric.data import Batch


# ── 高分辨率 FEM 参考解（保留作为回退） ──────────────────────────
def _compute_fem_reference(env):
    from skfem import Basis, ElementTriP1

    fp = env.fem_problem.fem_problem
    ref_mesh = fp._domain.get_integration_mesh()
    ref_basis = Basis(ref_mesh, ElementTriP1())
    ref_sol = fp.calculate_solution(ref_basis, cache=False)
    return ref_mesh, ref_sol


# ── Mie 解析参考解 ──────────────────────────────────────────────
def _compute_mie_reference(env):
    """Return Mie scattered field sampled at FEM mesh vertices.

    Falls back to FEM reference if scatterer is non-circular.
    """
    from environment.mie_solution import mie_scattered_field

    fp = getattr(env.fem_problem, "fem_problem", None)
    if fp is None:
        return _compute_fem_reference(env), None

    _eps_r = getattr(fp, "_eps_r", None)
    _radius = getattr(fp, "_radius", None)
    _cx = getattr(fp, "_cx", None)
    _cy = getattr(fp, "_cy", None)
    _k = getattr(fp, "_k", None)

    if any(v is None for v in [_eps_r, _radius, _cx, _cy, _k]):
        return _compute_fem_reference(env), None

    pts = env.mesh.p.T
    u_mie = mie_scattered_field(pts, k0=_k, eps_r=_eps_r, radius=_radius, cx=_cx, cy=_cy)

    from environment.mie_solution import mie_grid_solution
    import matplotlib.tri as tri

    xlim = (pts[:, 0].min(), pts[:, 0].max())
    ylim = (pts[:, 1].min(), pts[:, 1].max())
    grid = mie_grid_solution(_k, _eps_r, _radius, _cx, _cy,
                             x_range=xlim, y_range=ylim, Nx=500, Ny=500)

    mie_info = {
        "grid": grid,
        "eps_r": _eps_r, "radius": _radius,
        "cx": _cx, "cy": _cy, "k": _k,
    }
    return u_mie, mie_info


# ── 渲染辅助 ─────────────────────────────────────────────────────
def _render_field(ax, x, y, triang, values, title, vmin, vmax, show_mesh=True, cmap="jet"):
    tcf = ax.tripcolor(triang, values, shading="gouraud", cmap=cmap, vmin=vmin, vmax=vmax)
    if show_mesh and triang is not None:
        n = triang.triangles.shape[0]
        ax.triplot(triang, lw=(0.5 if n < 500 else 0.3), color="black",
                   alpha=(0.7 if n < 2000 else 0.5))
    ax.set_xlim(x.min(), x.max())
    ax.set_ylim(y.min(), y.max())
    ax.set_aspect("equal")
    ax.set_title(title, fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
    return tcf


# ── 保存 PNG ─────────────────────────────────────────────────────
def _save_png(steps, stem, checkpoint_path, k, cx=0.5, cy=0.5, radius=0.2, eps_r=2.0,
              mie_info=None):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import matplotlib.tri as tri

    per_step_dir = f"{stem}_steps"
    os.makedirs(os.path.dirname(stem) or ".", exist_ok=True)
    os.makedirs(per_step_dir, exist_ok=True)

    n = len(steps)
    ncols = min(n, 4)
    nrows = (n + ncols - 1) // ncols
    fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 3.5 * nrows))
    if nrows * ncols == 1:
        axes = np.array([axes])
    else:
        axes = np.array(axes).flatten()

    for i, step_data in enumerate(steps):
        mesh, scalar, err_val, n_elem = step_data[:4]
        pts = mesh.p.T
        tg = tri.Triangulation(pts[:, 0], pts[:, 1], mesh.t.T)
        s = np.abs(scalar) if np.iscomplexobj(scalar) else scalar
        lmin, lmax = s.min(), s.max()
        if lmax - lmin < 1e-12:
            lmin, lmax = lmin - 0.5, lmax + 0.5
        tcf = _render_field(axes[i], pts[:, 0], pts[:, 1], tg, s,
                            f"Step {i}: {n_elem} elem, err={err_val:.4f}",
                            lmin, lmax, cmap="jet")
        fig.colorbar(tcf, ax=axes[i], fraction=0.046, pad=0.04)
        axes[i].add_patch(plt.Circle((cx, cy), radius, fill=False,
                                     edgecolor="cyan", linewidth=1.5, linestyle="--"))

    for j in range(n, len(axes)):
        axes[j].set_visible(False)

    fig.subplots_adjust(left=0.04, right=0.90, top=0.90, bottom=0.06, wspace=0.15, hspace=0.30)
    k_str = f"k={k:.1f}" if k is not None else "k=?"
    ref_tag = " [Mie ref]" if mie_info is not None else ""
    fig.suptitle(
        f"Helmholtz |E_scat|{ref_tag} — {checkpoint_path}\n"
        f"{k_str}, eps_r={eps_r:.1f} at ({cx:.2f},{cy:.2f}) r={radius:.2f}",
        fontsize=12,
    )
    fig.savefig(f"{stem}.png", dpi=200, bbox_inches="tight")
    plt.close(fig)
    print(f"[Viz] Overview → {stem}.png")

    for i, step_data in enumerate(steps):
        mesh, scalar, err_val, n_elem = step_data[:4]
        u_mie_at_verts = step_data[4] if len(step_data) > 4 else None

        pts = mesh.p.T
        tg_coarse = tri.Triangulation(pts[:, 0], pts[:, 1], mesh.t.T)
        coarse_val = np.abs(scalar) if np.iscomplexobj(scalar) else scalar

        has_mie = u_mie_at_verts is not None
        ncols = 3 if has_mie else 1
        fig2, axes2 = plt.subplots(1, ncols, figsize=(6 * ncols, 6))
        axes2 = [axes2] if ncols == 1 else list(np.atleast_1d(axes2))

        # ── Panel 1: FEM scattered field ──
        cvmin, cvmax = coarse_val.min(), coarse_val.max()
        if cvmax - cvmin < 1e-12:
            cvmin, cvmax = cvmin - 0.5, cvmax + 0.5
        tcf1 = _render_field(axes2[0], pts[:, 0], pts[:, 1], tg_coarse, coarse_val,
                             f"Step {i}: FEM |E_scat| ({n_elem} elem)  max={cvmax:.4f}",
                             cvmin, cvmax, cmap="jet")
        axes2[0].add_patch(plt.Circle((cx, cy), radius, fill=False,
                                      edgecolor="cyan", linewidth=1.5, linestyle="--"))
        fig2.colorbar(tcf1, ax=axes2[0], fraction=0.046, pad=0.04)

        im2 = None
        if has_mie:
            # ── Panel 2: Mie scattered field (smooth grid, not FEM vertices) ──
            if mie_info is not None and "grid" in mie_info:
                g = mie_info["grid"]
                gm = np.abs(g["E_scat"])
                mvmin, mvmax = gm.min(), gm.max()
                if mvmax - mvmin < 1e-12:
                    mvmin, mvmax = mvmin - 0.5, mvmax + 0.5
                im2 = axes2[1].pcolormesh(g["X"], g["Y"], gm,
                                          shading="gouraud", cmap="jet",
                                          vmin=mvmin, vmax=mvmax)
                axes2[1].set_title(f"Mie |E_scat|  max={mvmax:.4f}", fontsize=9)
            else:
                mie_abs = np.abs(u_mie_at_verts)
                mvmin, mvmax = mie_abs.min(), mie_abs.max()
                if mvmax - mvmin < 1e-12:
                    mvmin, mvmax = mvmin - 0.5, mvmax + 0.5
                im2 = _render_field(axes2[1], pts[:, 0], pts[:, 1], tg_coarse, mie_abs,
                                    f"Mie |E_scat|  max={mvmax:.4f}",
                                    mvmin, mvmax, show_mesh=False, cmap="jet")
            axes2[1].set_aspect("equal")
            axes2[1].set_xticks([])
            axes2[1].set_yticks([])
            axes2[1].add_patch(plt.Circle((cx, cy), radius, fill=False,
                                          edgecolor="cyan", linewidth=1.5, linestyle="--"))
            if im2 is not None:
                fig2.colorbar(im2, ax=axes2[1], fraction=0.046, pad=0.04)

            # ── Panel 3: ||FEM| - |Mie|| error ──
            mie_abs = np.abs(u_mie_at_verts)
            error_abs = np.abs(coarse_val - mie_abs)
            evmin, evmax = 0.0, error_abs.max() or 1.0
            if evmax - evmin < 1e-12:
                evmax = evmin + 1.0
            tcf3 = _render_field(axes2[2], pts[:, 0], pts[:, 1], tg_coarse, error_abs,
                                 f"||FEM|-|Mie||  L2={err_val:.4f}  max={error_abs.max():.4f}",
                                 evmin, evmax, show_mesh=True, cmap="hot")
            axes2[2].add_patch(plt.Circle((cx, cy), radius, fill=False,
                                          edgecolor="cyan", linewidth=1.5, linestyle="--"))
            fig2.colorbar(tcf3, ax=axes2[2], fraction=0.046, pad=0.04)

        fig2.tight_layout()
        fig2.savefig(f"{per_step_dir}/step{i:02d}.png", dpi=150, bbox_inches="tight")
        plt.close(fig2)

    print(f"[Viz] Per-step PNGs → {per_step_dir}/  ({n} files)")


# ── Viz 模式入口 ──────────────────────────────────────────────────
def visualize(config: dict, checkpoint_path: str, output_path: str = "result/visualization.png",
              k_test=None, center=None, radius=None, eps_test=None):
    from src.network import create_model
    from src.utils import load_checkpoint, setup_helmholtz_config

    k = setup_helmholtz_config(config, k_test=k_test, center=center, radius=radius,
                               eps_test=eps_test)
    algo = config.get("algorithm", {})

    from environment.mesh_refinement import MeshRefinement

    env = MeshRefinement(
        environment_config=config.get("environment", {}).get("mesh_refinement", {}),
        seed=99,
    )
    model = create_model(env, config.get("network", {}), algo.get("ppo", {}))
    load_checkpoint(model, checkpoint_path)
    model.eval()

    stem = output_path.rsplit(".", 1)[0] if "." in output_path else output_path

    print(f"\n[Viz] Initializing...")
    obs = env.reset()

    _fp = getattr(env.fem_problem, "fem_problem", None)
    _cx = getattr(_fp, "_cx", 0.5) if _fp is not None else 0.5
    _cy = getattr(_fp, "_cy", 0.5) if _fp is not None else 0.5
    _radius = getattr(_fp, "_radius", 0.2) if _fp is not None else 0.2
    _eps_r = getattr(_fp, "_eps_r", 2.0) if _fp is not None else 2.0

    print(f"[Viz] Helmholtz params: k={k:.3f}  eps_r={_eps_r:.2f}  "
          f"center=({_cx:.3f}, {_cy:.3f})  radius={_radius:.3f}")

    # ── Mie analytical reference ──
    print(f"[Viz] Computing Mie reference solution...")
    u_mie_ref, mie_info = _compute_mie_reference(env)
    if mie_info is not None:
        print(f"[Viz] Mie reference ready (analytical, no domain truncation error)")

    # ── Initial step ──
    init_mesh = env.mesh
    init_sol = env.scalar_solution
    init_err = _compute_step_error(env, u_mie_ref)
    steps = [(init_mesh, init_sol, init_err, env.num_agents, u_mie_ref)]

    print(f"[Viz] Running inference...")
    done = False
    step_idx = 0
    while not done:
        with torch.no_grad():
            actions, _, _ = model(Batch.from_data_list([obs]), deterministic=True)
        obs, _, done, _ = env.step(actions.cpu().numpy())
        step_idx += 1
        sol = env.scalar_solution
        n_elem = env.num_agents
        u_mie_current = _eval_mie_on_mesh(env, mie_info)
        step_err = _compute_step_error(env, u_mie_current)

        diag_n_sel = getattr(env, "_diag_selected_count", -1)
        diag_n_elig = int(getattr(env, "_diag_eligible_ratio", 0) * env.num_agents)
        diag_n_mask = int(getattr(env, "_diag_masked_ratio", 0) * env.num_agents)
        remaining = getattr(env, "_n_budget", 0) - env.num_agents
        print(f"  Step {step_idx}: verts={env.mesh.p.shape[1]} elem={n_elem} "
              f"mie_err={step_err:.4f} "
              f"sel={diag_n_sel} elig={diag_n_elig} masked={diag_n_mask} "
              f"remaining={remaining} done={done}")

        steps.append((env.mesh, sol, step_err, n_elem, u_mie_current))

    _save_png(steps, stem, checkpoint_path, k, cx=_cx, cy=_cy, radius=_radius,
              eps_r=_eps_r, mie_info=mie_info)
    print(f"[Viz] Done → {output_path}")


def _compute_step_error(env, u_mie_ref) -> float:
    """相对 L₂ 误差: ||u_fem − u_mie||₂ / ||u_mie||₂  (复数，含幅值+相位)。"""
    if u_mie_ref is None:
        return float("nan")
    u_fem = env.scalar_solution       # complex scattered field
    diff = np.abs(u_fem - u_mie_ref)  # pointwise |complex difference|
    denom = np.linalg.norm(np.abs(u_mie_ref))
    if denom < 1e-12:
        denom = 1.0
    return float(np.linalg.norm(diff) / denom)


def _eval_mie_on_mesh(env, mie_info):
    """Re-evaluate Mie scattered field on current FEM mesh vertices."""
    if mie_info is None:
        return None
    from environment.mie_solution import mie_scattered_field
    pts = env.mesh.p.T
    return mie_scattered_field(pts, k0=mie_info["k"], eps_r=mie_info["eps_r"],
                               radius=mie_info["radius"], cx=mie_info["cx"], cy=mie_info["cy"])
