cols = [
    "mu",
    "ea", "ga22", "ga33", "gj", "ei22", "ei33",
    "stf12r", "stf13r", "stf14r", "stf15r", "stf16r",
    "stf23r", "stf24r", "stf25r", "stf26r",
    "stf34r", "stf35r", "stf36r",
    "stf45r", "stf46r",
    "stf56r",
    # "cmp12r", "cmp13r", "cmp14r", "cmp15r", "cmp16r",
    # "cmp23r", "cmp24r", "cmp25r", "cmp26r",
    # "cmp34r", "cmp35r", "cmp36r",
    # "cmp45r", "cmp46r",
    # "cmp56r",
    "tc2", "tc3", "sc2", "sc3"
    ]
log_cols = {"ea", "ga22", "ga33", "gj", "ei22", "ei33"}
stf_cols = {c for c in cols if c.startswith("stf")}
# cmp_cols = {c for c in cols if c.startswith("cmp")}
axis_names = {
    "mu":   "Mass per unit length",
    "ea":   "EA",
    "ga22": "GA_22",
    "ga33": "GA_33",
    "gj":   "GJ",
    "ei22": "EI_22",
    "ei33": "EI_33",
    "stf12r": "C_12",
    "stf13r": "C_13",
    "stf14r": "C_14",
    "stf15r": "C_15",
    "stf16r": "C_16",
    "stf23r": "C_23",
    "stf24r": "C_24",
    "stf25r": "C_25",
    "stf26r": "C_26",
    "stf34r": "C_34",
    "stf35r": "C_35",
    "stf36r": "C_36",
    "stf45r": "C_45",
    "stf46r": "C_46",
    "stf56r": "C_56",
    # "cmp12r": "S_12",
    # "cmp13r": "S_13",
    # "cmp23r": "S_23",
    # "cmp14r": "S_14",
    # "cmp15r": "S_15",
    # "cmp16r": "S_16",
    # "cmp24r": "S_24",
    # "cmp25r": "S_25",
    # "cmp26r": "S_26",
    # "cmp34r": "S_34",
    # "cmp35r": "S_35",
    # "cmp36r": "S_36",
    # "cmp45r": "S_45",
    # "cmp46r": "S_46",
    # "cmp56r": "S_56",
    "tc2":  "Tension center x_2",
    "tc3":  "Tension center x_3",
    "sc2":  "Shear center x_2",
    "sc3":  "Shear center x_3",
}
axis_units = {
    "mu":   "kg/m",
    "ea":   "N",
    "ga22": "N",
    "ga33": "N",
    "gj":   "N-m^2",
    "ei22": "N-m^2",
    "ei33": "N-m^2",
    "stf12r": "N",
    "stf13r": "N",
    "stf23r": "N",
    "stf14r": "N-m",
    "stf15r": "N-m",
    "stf16r": "N-m",
    "stf24r": "N-m",
    "stf25r": "N-m",
    "stf26r": "N-m",
    "stf34r": "N-m",
    "stf35r": "N-m",
    "stf36r": "N-m",
    "stf45r": "N-m^2",
    "stf46r": "N-m^2",
    "stf56r": "N-m^2",
    # "cmp12r": "1/N",
    # "cmp13r": "1/N",
    # "cmp23r": "1/N",
    # "cmp14r": "1/(N-m)",
    # "cmp15r": "1/(N-m)",
    # "cmp16r": "1/(N-m)",
    # "cmp24r": "1/(N-m)",
    # "cmp25r": "1/(N-m)",
    # "cmp26r": "1/(N-m)",
    # "cmp34r": "1/(N-m)",
    # "cmp35r": "1/(N-m)",
    # "cmp36r": "1/(N-m)",
    # "cmp45r": "1/(N-m^2)",
    # "cmp46r": "1/(N-m^2)",
    # "cmp56r": "1/(N-m^2)",
    "tc2":  "m",
    "tc3":  "m",
    "sc2":  "m",
    "sc3":  "m",
}
def title_of(c):
    return f"{axis_names.get(c, c.upper())} [{axis_units.get(c, '')}]"

def scale_of(c):
    if c in log_cols:
        return "log"
    if c in stf_cols:
        return "symlog"
    # if c in cmp_cols:
    #     return "symlog"
    return "linear"

def symlog_linthresh(c):
    values = np.concatenate([df[c].to_numpy(dtype=float) for df in datasets.values()])
    values = np.abs(values[np.isfinite(values)])
    values = values[values > 0]
    if values.size == 0:
        return 1.0
    return float(10 ** np.floor(np.log10(np.quantile(values, 0.1))))

def symlog_transform(values, linthresh):
    values = np.asarray(values, dtype=float)
    return np.sign(values) * np.log10(1 + np.abs(values) / linthresh)

def format_tick(v):
    if v == 0:
        return "0"
    if 1e-3 <= abs(v) < 1e3:
        return f"{v:g}"
    return f"{v:.0e}".replace("e+0", "e").replace("e-0", "e-")

def plot_values(df, c):
    values = df[c].to_numpy(dtype=float)
    if c in stf_cols:
        return symlog_transform(values, symlog_linthresh(c))
    # if c in cmp_cols:
    #     return symlog_transform(values, symlog_linthresh(c))
    return values

def axis_layout_for(c):
    if c in log_cols:
        return {"type": "log", "tickmode": "auto", "tickvals": None, "ticktext": None}
    if c in stf_cols:
    # if c in cmp_cols:
        # Plotly has no native symlog axis, so cmp* columns are drawn on a linear axis with symlog-transformed values and custom ticks.
        linthresh = symlog_linthresh(c)
        values = np.concatenate([df[c].to_numpy(dtype=float) for df in datasets.values()])
        values = np.abs(values[np.isfinite(values)])
        max_abs = float(values.max()) if values.size else linthresh
        min_exp = int(np.floor(np.log10(linthresh)))
        max_exp = int(np.ceil(np.log10(max_abs))) if max_abs > 0 else min_exp
        tick_values = {0.0}
        for exp in range(min_exp, max_exp + 1):
            tick = 10.0 ** exp
            if tick <= max_abs * 1.05:
                tick_values.update((-tick, tick))
        tick_values = sorted(tick_values)
        return {
            "type": "linear",
            "tickmode": "array",
            "tickvals": symlog_transform(tick_values, linthresh).tolist(),
            "ticktext": [format_tick(v) for v in tick_values],
        }
    return {"type": "linear", "tickmode": "auto", "tickvals": None, "ticktext": None}

def axis_update_args(axis, c):
    layout_updates = {f"{axis}axis.title.text": title_of(c)}
    layout_updates.update({f"{axis}axis.{k}": v for k, v in axis_layout_for(c).items()})
    return [
        {axis: [plot_values(df, c) for df in datasets.values()]},
        layout_updates,
    ]

datasets = {
    "al_001":      df_skin_al_001,
    "cfrp_001":    df_skin_cfrp_001,
    "cfrp_001_45": df_skin_cfrp_001_45,
    # "cfrp_001_90": df_skin_cfrp_001_90,
}
dataset_names = {
    "al_001":      "Aluminum",
    "cfrp_001":    "CFRP (0°)",
    "cfrp_001_45": "CFRP (45°)",
    # "cfrp_001_90": "CFRP (90°)",
}
dataset_markers = {
    "al_001":      "circle-open",
    "cfrp_001":    "arrow-open",
    "cfrp_001_45": "arrow-open",
    # "cfrp_001_90": "arrow-open",
}
dataset_marker_angles = {
    "al_001": 0,
    "cfrp_001": 0,
    "cfrp_001_45": 45,
    # "cfrp_001_90": 90,
}
dataset_colors = {
    "al_001": "#636EFA",
    "cfrp_001": "#EF553B",
    "cfrp_001_45": "#00CC96",
    # "cfrp_001_90": "#AB63FA",
}

def filled_marker_symbol(symbol):
    return symbol[:-5] if symbol.endswith("-open") else symbol

def marker_style(name, highlighted=False):
    return dict(
        symbol=filled_marker_symbol(dataset_markers[name]) if highlighted else dataset_markers[name],
        size=10 if highlighted else 5,
        opacity=1.0 if highlighted else 0.2,
        color=dataset_colors[name],
        angle=-1 * dataset_marker_angles[name] + 90,  # rotate markers to match fiber angle
        line=dict(color='black', width=1),
    )

def highlight_mask(df, airfoil_name):
    if airfoil_name == "(none)":
        return np.zeros(len(df), dtype=bool)
    return df["airfoil_name"].astype(str).eq(airfoil_name).to_numpy()

airfoil_names = sorted(
    set().union(*[set(df["airfoil_name"].dropna().astype(str)) for df in datasets.values()])
)
airfoil_options = ["(none)"] + airfoil_names
x0, y0 = "ei22", "ei33"

axis_titles = {c: title_of(c) for c in cols}
axis_layouts = {c: axis_layout_for(c) for c in cols}
datasets_payload = {}

fig = go.Figure()
trace_indices = {}
for name, df in datasets.items():
    airfoil_text = df["airfoil_name"].astype(str).tolist()
    datasets_payload[name] = {
        "label": dataset_names[name],
        "airfoil_names": airfoil_text,
        "values": {c: plot_values(df, c).tolist() for c in cols},
    }
    fig.add_trace(
        go.Scatter(
            x=plot_values(df, x0),
            y=plot_values(df, y0),
            mode="markers",
            marker=marker_style(name, highlighted=False),
            name=dataset_names[name],
            legendgroup=name,
            text=airfoil_text,
            hovertemplate=f"<b>{dataset_names[name]}</b><br>airfoil: %{{text}}<br>x=%{{x}}<br>y=%{{y}}<extra></extra>",
        )
    )
    background_idx = len(fig.data) - 1
    fig.add_trace(
        go.Scatter(
            x=[],
            y=[],
            mode="markers",
            marker=marker_style(name, highlighted=True),
            name=f"{dataset_names[name]} highlight",
            legendgroup=name,
            showlegend=False,
            text=[],
            hovertemplate=f"<b>{dataset_names[name]}</b><br>airfoil: %{{text}}<br>x=%{{x}}<br>y=%{{y}}<extra></extra>",
        )
    )
    trace_indices[name] = {
        "background": background_idx,
        "highlight": len(fig.data) - 1,
    }

fig.update_layout(
    template="plotly_white",
    height=600,
    width=600,
    xaxis=dict(title=dict(text=title_of(x0)), **axis_layout_for(x0)),
    yaxis=dict(title=dict(text=title_of(y0)), **axis_layout_for(y0)),
    margin=dict(t=40),
    legend=dict(
        title=dict(text="Datasets<br><sub>(click to toggle)</sub>"),
        groupclick="togglegroup",
    ),
)

plot_div_id = "airfoil-cs-plot"
x_select_id = f"{plot_div_id}-x"
y_select_id = f"{plot_div_id}-y"
airfoil_select_id = f"{plot_div_id}-airfoil"

def option_html(options, selected_value, label_fn=lambda value: value):
    parts = []
    for value in options:
        selected_attr = " selected" if value == selected_value else ""
        parts.append(
            f'<option value="{html.escape(value)}"{selected_attr}>{html.escape(label_fn(value))}</option>'
        )
    return "".join(parts)

controls_html = f"""
<div style=\"display:flex;gap:12px;flex-wrap:wrap;align-items:center;margin:0 0 12px 0;\">
  <label>X:
    <select id=\"{x_select_id}\" style=\"margin-left:6px;\">{option_html(cols, x0, lambda c: axis_names.get(c, c.upper()))}</select>
  </label>
  <label>Y:
    <select id=\"{y_select_id}\" style=\"margin-left:6px;\">{option_html(cols, y0, lambda c: axis_names.get(c, c.upper()))}</select>
  </label>
  <label>Airfoil:
    <select id=\"{airfoil_select_id}\" style=\"margin-left:6px;max-width:220px;\">{option_html(airfoil_options, '(none)')}</select>
  </label>
</div>
"""

payload = {
    "dataset_order": list(datasets.keys()),
    "datasets": datasets_payload,
    "axis_titles": axis_titles,
    "axis_layouts": axis_layouts,
    "none_airfoil": "(none)",
}
fig_html = fig.to_html(full_html=False, include_plotlyjs="cdn", div_id=plot_div_id)
script = f"""
<script>
(function() {{
  const payload = {json.dumps(payload, separators=(',', ':'))};
  const graphDiv = document.getElementById('{plot_div_id}');
  const xSelect = document.getElementById('{x_select_id}');
  const ySelect = document.getElementById('{y_select_id}');
  const airfoilSelect = document.getElementById('{airfoil_select_id}');

  function highlightedPoints(dataset, xCol, yCol, selectedAirfoil) {{
    const x = [];
    const y = [];
    const text = [];
    if (selectedAirfoil === payload.none_airfoil) {{
      return {{ x, y, text }};
    }}

    const names = dataset.airfoil_names;
    const xValues = dataset.values[xCol];
    const yValues = dataset.values[yCol];
    for (let i = 0; i < names.length; i += 1) {{
      if (names[i] === selectedAirfoil) {{
        x.push(xValues[i]);
        y.push(yValues[i]);
        text.push(names[i]);
      }}
    }}
    return {{ x, y, text }};
  }}

  function updatePlot() {{
    const xCol = xSelect.value;
    const yCol = ySelect.value;
    const selectedAirfoil = airfoilSelect.value;
    const update = {{ x: [], y: [], text: [] }};

    payload.dataset_order.forEach((name) => {{
      const dataset = payload.datasets[name];
      update.x.push(dataset.values[xCol]);
      update.y.push(dataset.values[yCol]);
      update.text.push(dataset.airfoil_names);

      const highlight = highlightedPoints(dataset, xCol, yCol, selectedAirfoil);
      update.x.push(highlight.x);
      update.y.push(highlight.y);
      update.text.push(highlight.text);
    }});

    Plotly.update(graphDiv, update, {{
      xaxis: Object.assign({{ title: {{ text: payload.axis_titles[xCol] }} }}, payload.axis_layouts[xCol]),
      yaxis: Object.assign({{ title: {{ text: payload.axis_titles[yCol] }} }}, payload.axis_layouts[yCol]),
    }});
  }}

  xSelect.addEventListener('change', updatePlot);
  ySelect.addEventListener('change', updatePlot);
  airfoilSelect.addEventListener('change', updatePlot);
}})();
</script>
"""

html_output_path = "_static/plot.interactive.html"
html_document = f"""<!doctype html>
<html lang=\"en\">
<head>
  <meta charset=\"utf-8\" />
  <meta name=\"viewport\" content=\"width=device-width, initial-scale=1\" />
  <title>Airfoil cross-section plot</title>
  <style>
    body {{ margin: 0; padding: 16px; font-family: sans-serif; }}
  </style>
</head>
<body>
{controls_html}
{fig_html}
{script}
</body>
</html>
"""
with open(html_output_path, "w", encoding="utf-8") as f:
    f.write(html_document)

display(HTML(controls_html + fig_html + script))
# display(HTML(fig.to_html(full_html=False, include_plotlyjs="cdn", include_mathjax="cdn")))