from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import plotly.graph_objects as go
from plotly.basedatatypes import BaseTraceType
from skfem import Mesh


# 将网格与标量场转为 plotly 三角形 traces + 布局，供 RL 环境实时渲染
def get_plotly_mesh_traces_and_layout(
    mesh: Mesh,
    scalars: np.ndarray,
    title: str = "Mesh",
    mesh_dimension: int = 2,
    boundary: Optional[np.ndarray] = None,
) -> Tuple[List[BaseTraceType], Dict[str, Any]]:
    vertices = mesh.p
    triangles = mesh.t
    n_elements = triangles.shape[1]
    s = np.asarray(scalars, dtype=np.float64).flatten()

    x_tri = vertices[0, triangles].T
    y_tri = vertices[1, triangles].T
    intensity_tri = s[triangles].T

    vmin, vmax = s.min(), s.max()

    traces = []
    for elem_idx in range(n_elements):
        x_e, y_e, s_e = x_tri[elem_idx], y_tri[elem_idx], intensity_tri[elem_idx]
        traces.append(go.Scatter(
            x=x_e.tolist() + [x_e[0]],
            y=y_e.tolist() + [y_e[0]],
            mode="lines",
            fill="toself",
            fillcolor=_get_color(float(np.mean(s_e)), vmin, vmax),
            line=dict(color="black", width=0.5),
            showlegend=False,
            hoverinfo="skip",
        ))

    if traces:
        traces[0].marker = dict(
            color=s.min(), colorscale="RdBu_r", showscale=True,
            colorbar=dict(title="Solution"),
        )

    layout = {
        "title": title,
        "xaxis": {"title": "x", "scaleanchor": "y"},
        "yaxis": {"title": "y"},
        "showlegend": False,
    }
    if boundary is not None:
        layout["xaxis"]["range"] = [boundary[0], boundary[2]]
        layout["yaxis"]["range"] = [boundary[1], boundary[3]]

    return traces, layout


# 标量值 → matplotlib RdBu_r 色表映射的 RGBA 字符串
def _get_color(value: float, vmin: float, vmax: float) -> str:
    import matplotlib.cm as cm
    import matplotlib.colors as mcolors

    norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
    rgba = cm.RdBu_r(norm(value))
    r, g, b, a = rgba
    return f"rgba({int(r * 255)},{int(g * 255)},{int(b * 255)},{a:.2f})"