import io
import json
from typing import Dict, Tuple
from itertools import combinations

import imageio.v2 as imageio
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import streamlit as st
import networkx as nx

from ngf import NGFGenerator, EnergyDistribution

st.set_page_config(page_title="NGF Generator (Bianconi–Rahmede) + GIF", layout="wide")

# ---------- Sidebar Controls ----------
st.sidebar.title("NGF Controls")

with st.sidebar:
    st.markdown("**Network Geometry with Flavor (NGF)**")
    d = st.number_input("Dimension d", min_value=1, max_value=3, value=2, step=1)
    s = st.selectbox("Flavor s", options=[-1, 0, 1], index=1)
    beta = st.number_input("Inverse temperature β", min_value=0.0, value=0.1, step=0.1)

    if "total_nodes_exact" not in st.session_state:
        st.session_state.total_nodes_exact = 1000
    st.caption("N (exact)")
    total_nodes = st.number_input(
        "N (exact)",
        min_value=2, max_value=50000, value=int(st.session_state.total_nodes_exact), step=1,
        label_visibility="collapsed",
        help="Exakten Knotenwert eingeben (2–50000)."
    )
    st.session_state.total_nodes_exact = int(total_nodes)

    seed = st.number_input("Random seed", min_value=0, value=42, step=1)

    st.divider()
    st.markdown("**Energy distribution g(ε)**")
    dist_kind = st.selectbox("Kind", ["uniform_int", "poisson", "custom_list"], index=0)
    if dist_kind == "uniform_int":
        e_low = st.number_input("ε low (inclusive)", value=0, step=1)
        e_high = st.number_input("ε high (exclusive)", value=10, step=1)
        dist_params = {"low": int(e_low), "high": int(e_high)}
    elif dist_kind == "poisson":
        lam = st.number_input("λ (Poisson mean)", value=4.0, step=0.5)
        dist_params = {"lam": float(lam)}
    else:
        vals = st.text_input("Custom integer values (comma-separated)", "0,1,2,3,5,8,13")
        try:
            arr = [int(x.strip()) for x in vals.split(",") if x.strip()!=""]
        except Exception:
            arr = [0,1,2,3]
        dist_params = {"values": arr}

    st.divider()
    st.markdown("**Visualization options**")
    show_triangles = st.checkbox("Show 2-faces (triangles) as mesh (if d≥2)", value=True)
    show_boundary = st.checkbox("Highlight boundary (d-1)-faces", value=False)
    size_by_degree = st.checkbox("Scale node size by 1-skeleton degree (Fig. 2 style)", value=True)

    st.divider()
    do_generate = st.button("Generate NGF", type="primary")

    # Animation controls under the main button
    fps = st.number_input("Animation FPS", min_value=1, max_value=60, value=12, step=1)
    step_N = st.number_input("ΔN pro Frame", min_value=1, max_value=5000, value=20, step=1)
    do_anim = st.button("Generate Animation")


@st.cache_data(show_spinner=False)
def build_ngf(d, s, beta, total_nodes, dist_kind, dist_params, seed):
    gen = NGFGenerator(
        d=int(d),
        s=int(s),
        beta=float(beta),
        energy_dist=EnergyDistribution(kind=dist_kind, params=dist_params, seed=seed),
        seed=int(seed),
    )
    gen.run(total_nodes=int(total_nodes))
    return gen


def graph_layout_3d(G: nx.Graph, seed: int = 42) -> Dict[int, Tuple[float, float, float]]:
    """Preferred: spring_layout in 3D (requires SciPy). Fallback: random_layout (SciPy-free)."""
    try:
        import scipy  # probe availability
        pos = nx.spring_layout(G, dim=3, seed=seed, k=None, iterations=200)
    except Exception:
        pos = nx.random_layout(G, center=(0.0, 0.0, 0.0), dim=3, seed=seed)
    return {int(k): (float(v[0]), float(v[1]), float(v[2])) for k, v in pos.items()}


def plot_ngf_3d(gen: NGFGenerator, show_triangles: bool = True, show_boundary: bool = False, size_by_degree: bool = True):
    from itertools import combinations
    G = gen.to_networkx()
    pos = graph_layout_3d(G, seed=123)

    import numpy as np
    energies = np.array([gen.node_energy[i] for i in gen.nodes], dtype=float)
    deg1 = gen.node_degrees_1skeleton()
    degrees = np.array([deg1[i] for i in gen.nodes], dtype=float)

    # Node size ∝ degree (robust scaling)
    if size_by_degree and len(degrees):
        deg_min = degrees.min()
        deg_range = float(np.ptp(degrees))
        size_min, size_max = 6.0, 22.0
        if deg_range <= 1e-12:
            node_sizes = np.full_like(energies, (size_min + size_max) * 0.5)
        else:
            node_sizes = size_min + (size_max - size_min) * (degrees - deg_min) / deg_range
    else:
        node_sizes = np.full_like(energies, 10.0)

    xs = [pos[i][0] for i in gen.nodes]
    ys = [pos[i][1] for i in gen.nodes]
    zs = [pos[i][2] for i in gen.nodes]

    import plotly.graph_objects as go
    node_trace = go.Scatter3d(
        x=xs, y=ys, z=zs,
        mode="markers",
        marker=dict(
            size=node_sizes,
            color=energies if len(energies) else [0.0],
            cmin=float(energies.min()) if len(energies) else 0.0,
            cmax=float(energies.max()) if len(energies) else 1.0,
            colorscale=[
                [0.0, "#ff0000"], [0.2, "#ff7f00"], [0.4, "#ffff00"],
                [0.6, "#00ff00"], [0.8, "#0000ff"], [1.0, "#8b00ff"],
            ],
            showscale=True,
            colorbar=dict(title="ε"),
            opacity=0.9,
        ),
        text=[f"node {i}<br>ε={gen.node_energy[i]}<br>deg={deg1[i]}" for i in gen.nodes],
        hoverinfo="text",
        name="nodes",
    )

    # Edges
    ex, ey, ez = [], [], []
    for u, v in gen.edges:
        x0, y0, z0 = pos[u]
        x1, y1, z1 = pos[v]
        ex += [x0, x1, None]
        ey += [y0, y1, None]
        ez += [z0, z1, None]
    edge_trace = go.Scatter3d(
        x=ex, y=ey, z=ez,
        mode="lines",
        line=dict(width=2.0, color="rgba(90,90,90,0.75)"),
        hoverinfo="none",
        name="edges",
        opacity=0.75,
    )

    data = [edge_trace, node_trace]

    # Triangles (2-faces)
    if show_triangles and gen.d >= 2:
        tris_set = set()
        if gen.d == 2:
            for s in gen.simplexes_d:
                if len(s) >= 3:
                    a, b, c = sorted(list(s))[:3]
                    tris_set.add((a, b, c))
        else:
            for s in gen.simplexes_d:
                for (a, b, c) in combinations(s, 3):
                    tris_set.add(tuple(sorted((a, b, c))))

        if tris_set:
            idx_map = {nid: idx for idx, nid in enumerate(gen.nodes)}
            I, J, K = [], [], []
            for (a, b, c) in tris_set:
                if a in idx_map and b in idx_map and c in idx_map:
                    I.append(idx_map[a]); J.append(idx_map[b]); K.append(idx_map[c])
            if I:
                mesh = go.Mesh3d(
                    x=xs, y=ys, z=zs,
                    i=I, j=J, k=K,
                    opacity=0.18,
                    name="2-faces (triangles)",
                    hoverinfo="skip"
                )
                data.insert(0, mesh)

    if show_boundary and gen.d >= 1:
        bx, by, bz = [], [], []
        boundary = gen.get_boundary_faces()
        for face in boundary:
            nodes = list(face)
            for i in range(len(nodes)):
                for j in range(i + 1, len(nodes)):
                    u, v = nodes[i], nodes[j]
                    x0, y0, z0 = pos[u]
                    x1, y1, z1 = pos[v]
                    bx += [x0, x1, None]
                    by += [y0, y1, None]
                    bz += [z0, z1, None]
        if bx:
            boundary_trace = go.Scatter3d(
                x=bx, y=by, z=bz, mode="lines",
                line=dict(width=4, color="rgba(50,50,50,0.35)"),
                opacity=0.3, name="boundary (d-1)-faces", hoverinfo="skip"
            )
            data.append(boundary_trace)

    fig = go.Figure(data=data)
    fig.update_layout(
        scene=dict(xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False)),
        margin=dict(l=0, r=0, t=30, b=0),
        legend=dict(orientation="h"),
        height=700,
        title=f"NGF d={gen.d}, s={gen.s}, β={gen.beta}, N={gen.N}, V={len(gen.simplexes_d)} (Fig. 2 style)"
    )
    return fig


def render_frame_png_matplotlib(gen, pos, uptoN, show_triangles, show_boundary, size_by_degree,
                                width=1200, height=800, dpi=150):
    """Matplotlib-Frame ohne Kaleido, mit skalierter Knotengröße (Grad) relativ zur GIF-Größe."""
    import io
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.colors import LinearSegmentedColormap, Normalize
    from mpl_toolkits.mplot3d.art3d import Poly3DCollection
    from itertools import combinations
    import numpy as np

    nodes_sub = [i for i in gen.nodes if i < uptoN]
    node_set = set(nodes_sub)
    edges_sub = [(u, v) for (u, v) in gen.edges if u in node_set and v in node_set]

    energies = np.array([gen.node_energy[i] for i in nodes_sub], dtype=float)

    # Grad im Teilgraph
    deg = {i: 0 for i in nodes_sub}
    for u, v in edges_sub:
        deg[u] += 1
        deg[v] += 1
    degrees = np.array([deg[i] for i in nodes_sub], dtype=float) if nodes_sub else np.array([])

    # --- Knoten-Skalierung: Pixel -> Points² (Matplotlib erwartet Fläche in pt²)
    min_side = float(min(width, height))
    # sinnvolle Relationen zur Bildgröße (Passe bei Bedarf an)
    min_px = max(4.0, 0.008 * min_side)   # ~6–10 px bei 800–1200
    max_px = max(10.0, 0.030 * min_side)  # ~24–36 px bei 800–1200

    if size_by_degree and degrees.size:
        dmin = float(degrees.min())
        drng = float(np.ptp(degrees))
        if drng <= 1e-12:
            # alle Grade (fast) gleich -> uniforme Größe
            sizes_px = np.full_like(degrees, (min_px + max_px) * 0.5)
        else:
            dnorm = (degrees - dmin) / drng
            sizes_px = min_px + (max_px - min_px) * dnorm
    else:
        sizes_px = np.full_like(energies, (min_px + max_px) * 0.5)

    # Umrechnung Pixel -> Points²
    px_to_pt = 72.0 / float(dpi)
    sizes_pts2 = (sizes_px * px_to_pt) ** 2

    xs = np.array([pos[i][0] for i in nodes_sub]) if nodes_sub else np.array([])
    ys = np.array([pos[i][1] for i in nodes_sub]) if nodes_sub else np.array([])
    zs = np.array([pos[i][2] for i in nodes_sub]) if nodes_sub else np.array([])

    # Fig.2 Farbskala (Energie)
    cmap = LinearSegmentedColormap.from_list(
        "rainbow6", ["#ff0000", "#ff7f00", "#ffff00", "#00ff00", "#0000ff", "#8b00ff"]
    )
    if energies.size:
        norm = Normalize(vmin=float(energies.min()), vmax=float(energies.max()))
        colors = cmap(norm(energies))
    else:
        from numpy import array
        norm = Normalize(vmin=0.0, vmax=1.0)
        colors = cmap(norm(np.array([0.0])))

    fig = plt.figure(figsize=(width/100.0, height/100.0), dpi=dpi)
    ax = fig.add_subplot(111, projection="3d")
    ax.set_axis_off()

    # Dreiecks-Mesh (2-faces)
    if show_triangles and gen.d >= 2 and len(nodes_sub) >= 3:
        tris = []
        if gen.d == 2:
            for s in gen.simplexes_d:
                if max(s) < uptoN and len(s) >= 3:
                    a, b, c = sorted(list(s))[:3]
                    if a in node_set and b in node_set and c in node_set:
                        tris.append([(pos[a][0], pos[a][1], pos[a][2]),
                                     (pos[b][0], pos[b][1], pos[b][2]),
                                     (pos[c][0], pos[c][1], pos[c][2])])
        else:
            for s in gen.simplexes_d:
                if max(s) < uptoN:
                    for (a, b, c) in combinations(s, 3):
                        if a in node_set and b in node_set and c in node_set:
                            tris.append([(pos[a][0], pos[a][1], pos[a][2]),
                                         (pos[b][0], pos[b][1], pos[b][2]),
                                         (pos[c][0], pos[c][1], pos[c][2])])
        if tris:
            mesh = Poly3DCollection(tris, alpha=0.15, facecolor="#9aa0a6", edgecolor="none")
            ax.add_collection3d(mesh)

    # Kantenbreite relativ zur Bildgröße
    lw_edge = max(1.0, 0.0015 * min_side)
    for u, v in edges_sub:
        x0, y0, z0 = pos[u]; x1, y1, z1 = pos[v]
        ax.plot([x0, x1], [y0, y1], [z0, z1], linewidth=lw_edge, color=(0.35, 0.35, 0.35, 0.8))

    # Randkanten (optional etwas dicker)
    if show_boundary and gen.d >= 1:
        lw_b = max(1.8, 0.0028 * min_side)
        for face, n in gen.face_n.items():
            if n == 0 and all(v in node_set for v in face):
                f = list(face)
                for i in range(len(f)):
                    for j in range(i + 1, len(f)):
                        u, v = f[i], f[j]
                        x0, y0, z0 = pos[u]; x1, y1, z1 = pos[v]
                        ax.plot([x0, x1], [y0, y1], [z0, z1], linewidth=lw_b, color=(0.2, 0.2, 0.2, 0.35))

    # Knoten
    if xs.size:
        ax.scatter(xs, ys, zs, s=sizes_pts2, c=colors, depthshade=True)

        # halbwegs gleiche Achsenskalierung
        mins = np.array([xs.min(), ys.min(), zs.min()])
        maxs = np.array([xs.max(), ys.max(), zs.max()])
        centers = (mins + maxs) / 2.0
        span = (maxs - mins).max() * 0.55 + 1e-9
        ax.set_xlim(centers[0]-span, centers[0]+span)
        ax.set_ylim(centers[1]-span, centers[1]+span)
        ax.set_zlim(centers[2]-span, centers[2]+span)

    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=dpi, bbox_inches="tight", pad_inches=0.02)
    buf.seek(0)
    plt.close(fig)
    return buf.getvalue()



# ---------- App Body ----------
st.title("Network Geometry with Flavor — NGF Generator & 3D Viewer (+ GIF, no-kaleido, HF-ready)")
st.caption("Fig. 2 style: node color ~ energy (red→violet), node size ~ degree; 2-faces optional.")

if do_generate:
    with st.spinner("Generating NGF..."):
        gen = build_ngf(d, s, beta, total_nodes, dist_kind, dist_params, seed)

    t1, t2, t3, t4 = st.tabs(["3D View", "Distributions", "Growth", "Export"])

    with t1:
        fig = plot_ngf_3d(gen, show_triangles=show_triangles, show_boundary=show_boundary, size_by_degree=size_by_degree)
        st.plotly_chart(fig, use_container_width=True)

        col1, col2, col3, col4 = st.columns(4)
        with col1:
            st.metric("N (nodes)", f"{gen.N}")
        with col2:
            st.metric("V (d-simplices)", f"{len(gen.simplexes_d)}")
        with col3:
            st.metric("Area A (boundary faces)", f"{sum(1 for f,n in gen.face_n.items() if n==0)}")
        with col4:
            st.metric("Edges (1-skeleton)", f"{len(gen.edges)}")

    with t2:
        deg1 = gen.node_degrees_1skeleton()
        hist_deg = {}
        for k in deg1.values():
            hist_deg[k] = hist_deg.get(k, 0) + 1
        xs = np.array(list(hist_deg.keys()), dtype=float)
        ys = np.array(list(hist_deg.values()), dtype=float)
        fig_deg = go.Figure(); fig_deg.add_bar(x=xs, y=ys)
        fig_deg.update_layout(title="1-skeleton degree distribution", xaxis_title="k", yaxis_title="count", height=360)
        fig_deg.update_xaxes(type="log"); fig_deg.update_yaxes(type="log")
        st.plotly_chart(fig_deg, use_container_width=True)

        if gen.d >= 1:
            hist_d1 = gen.generalized_degree_distribution(gen.d - 1)
            xs = np.array(list(hist_d1.keys()), dtype=float); ys = np.array(list(hist_d1.values()), dtype=float)
            fig_d1 = go.Figure(); fig_d1.add_bar(x=xs, y=ys)
            fig_d1.update_layout(title=f"Generalized degrees k₍{gen.d},{gen.d-1}₎ of (d-1)-faces", height=360)
            fig_d1.update_xaxes(type="log"); fig_d1.update_yaxes(type="log")
            st.plotly_chart(fig_d1, use_container_width=True)

        hist_nodes = gen.generalized_degree_distribution(0)
        xs = np.array(list(hist_nodes.keys()), dtype=float); ys = np.array(list(hist_nodes.values()), dtype=float)
        fig_n = go.Figure(); fig_n.add_bar(x=xs, y=ys)
        fig_n.update_layout(title=f"Generalized degrees k₍{gen.d},0₎ of nodes", height=360)
        fig_n.update_xaxes(type="log"); fig_n.update_yaxes(type="log")
        st.plotly_chart(fig_n, use_container_width=True)

        if gen.d >= 2:
            hist_links = gen.generalized_degree_distribution(1)
            xs = np.array(list(hist_links.keys()), dtype=float); ys = np.array(list(hist_links.values()), dtype=float)
            fig_l = go.Figure(); fig_l.add_bar(x=xs, y=ys)
            fig_l.update_layout(title=f"Generalized degrees k₍{gen.d},1₎ of links", height=360)
            fig_l.update_xaxes(type="log"); fig_l.update_yaxes(type="log")
            st.plotly_chart(fig_l, use_container_width=True)

    with t4:
        data = gen.to_json()
        st.json({"dimension": data["dimension"], "flavor": data["flavor"], "beta": data["beta"], "N": data["N"], "V": len(gen.simplexes_d)})
        j = json.dumps(data).encode("utf-8")
        st.download_button("Download NGF as JSON", j, file_name=f"ngf_d{gen.d}_s{gen.s}_N{gen.N}.json", mime="application/json")

# --- Animation (matplotlib-based, BytesIO download) ---
if do_anim:
    with st.spinner("Baue Animation (ohne Kaleido, HF-ready)..."):
        gen_anim = build_ngf(d, s, beta, total_nodes, dist_kind, dist_params, seed)
        G_final = gen_anim.to_networkx()
        pos = graph_layout_3d(G_final, seed=123)

        startN = gen_anim.d + 1
        endN = int(total_nodes)
        steps = list(range(startN, endN + 1, int(step_N)))
        if steps[-1] != endN:
            steps.append(endN)

        max_frames = 600
        if len(steps) > max_frames:
            st.warning(f"Viele Frames ({len(steps)}). Kürze auf {max_frames}. Erhöhe ΔN.")
            import numpy as np
            keep = np.linspace(0, len(steps)-1, max_frames).astype(int).tolist()
            steps = [steps[i] for i in keep]

        frames = []
        prog = st.progress(0.0)
        for i, upto in enumerate(steps):
            png_bytes = render_frame_png_matplotlib(gen_anim, pos, upto,
                                                    show_triangles=show_triangles,
                                                    show_boundary=show_boundary,
                                                    size_by_degree=size_by_degree,
                                                    width=1200, height=800, dpi=150)
            frames.append(imageio.imread(io.BytesIO(png_bytes)))
            prog.progress((i + 1) / len(steps))

        # Write GIF to memory buffer for direct download (no filesystem path needed)
        gif_buf = io.BytesIO()
        imageio.mimsave(gif_buf, frames, format="GIF", fps=int(fps))
        gif_buf.seek(0)

        st.success(f"Animation erstellt: {len(steps)} Frames @ {int(fps)} FPS")
        st.download_button("Download GIF", data=gif_buf, file_name="ngf_animation.gif", mime="image/gif")