feat: simple analysis page with sociogram

analysis.py

@@ -0,0 +1,144 @@
+from datetime import datetime
+import numpy as np
+import io
+import base64
+from fastapi import APIRouter
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel, Field
+from sqlmodel import Session, func, select
+from sqlmodel.sql.expression import SelectOfScalar
+from db import Chemistry, Player, engine
+import networkx as nx
+
+import matplotlib
+
+matplotlib.use("agg")
+import matplotlib.pyplot as plt
+
+analysis_router = APIRouter(prefix="/analysis")
+
+
+C = Chemistry
+P = Player
+
+
+def sociogram_json():
+    nodes = []
+    necessary_nodes = set()
+    links = []
+    with Session(engine) as session:
+        for p in session.exec(select(P)).fetchall():
+            nodes.append({"id": p.name, "appearance": 1})
+        subquery = (
+            select(C.user, func.max(C.time).label("latest"))
+            .where(C.time > datetime(2025, 2, 1, 10))
+            .group_by(C.user)
+            .subquery()
+        )
+        statement2 = select(C).join(
+            subquery, (C.user == subquery.c.user) & (C.time == subquery.c.latest)
+        )
+        for c in session.exec(statement2):
+            # G.add_node(c.user)
+            necessary_nodes.add(c.user)
+            for p in c.love:
+                # G.add_edge(c.user, p)
+                # p_id = session.exec(select(P.id).where(P.name == p)).one()
+                necessary_nodes.add(p)
+                links.append({"source": c.user, "target": p})
+    # nodes = [n for n in nodes if n["name"] in necessary_nodes]
+    return JSONResponse({"nodes": nodes, "links": links})
+
+
+def sociogram_data():
+    nodes = []
+    links = []
+    G = nx.DiGraph()
+    with Session(engine) as session:
+        for p in session.exec(select(P)).fetchall():
+            nodes.append({"id": p.name})
+            G.add_node(p.name)
+        subquery = (
+            select(C.user, func.max(C.time).label("latest"))
+            .where(C.time > datetime(2025, 2, 1, 10))
+            .group_by(C.user)
+            .subquery()
+        )
+        statement2 = (
+            select(C)
+            .where(C.user.in_(["Kruse", "Franz", "ck"]))
+            .join(subquery, (C.user == subquery.c.user) & (C.time == subquery.c.latest))
+        )
+        for c in session.exec(statement2):
+            for p in c.love:
+                G.add_edge(c.user, p)
+                links.append({"source": c.user, "target": p})
+    return G
+
+
+class Params(BaseModel):
+    node_size: int | None = Field(default=2400, alias="nodeSize")
+    font_size: int | None = Field(default=10, alias="fontSize")
+    arrow_size: int | None = Field(default=20, alias="arrowSize")
+    edge_width: float | None = Field(default=1, alias="edgeWidth")
+
+
+def sociogram_image(params: Params):
+    print(params)
+    plt.figure(figsize=(16, 10), facecolor="none")
+    ax = plt.gca()
+    ax.set_facecolor("none")  # Set the axis face color to none (transparent)
+    ax.axis("off")  # Turn off axis ticks and frames
+
+    G = sociogram_data()
+    pos = nx.spring_layout(
+        G, scale=2, k=1 / np.sqrt(G.number_of_edges()), iterations=50, seed=42
+    )
+    nx.draw_networkx_nodes(
+        G,
+        pos,
+        node_color="#99ccff",
+        edgecolors="#404040",
+        linewidths=1,
+        node_size=params.node_size,
+        alpha=0.86,
+    )
+    nx.draw_networkx_labels(G, pos, font_size=params.font_size)
+    nx.draw_networkx_edges(
+        G,
+        pos,
+        arrows=True,
+        edge_color="#404040",
+        arrowsize=params.arrow_size,
+        node_size=params.node_size,
+        width=params.edge_width,
+    )
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png", bbox_inches="tight", dpi=300, transparent=True)
+    buf.seek(0)
+    encoded_image = base64.b64encode(buf.read()).decode("UTF-8")
+    plt.close()
+
+    return {"image": encoded_image}
+
+
+analysis_router.add_api_route("/json", endpoint=sociogram_json, methods=["GET"])
+analysis_router.add_api_route("/image", endpoint=sociogram_image, methods=["POST"])
+
+if __name__ == "__main__":
+    with Session(engine) as session:
+        statement: SelectOfScalar[P] = select(func.count(P.id))
+        print("players in DB: ", session.exec(statement).first())
+    G = sociogram_data()
+    pos = nx.spring_layout(G, scale=1, k=2, iterations=50, seed=42)
+    edges = nx.draw_networkx_edges(
+        G,
+        pos,
+        arrows=True,
+        arrowsize=12,
+    )
+    nx.draw_networkx(
+        G, pos, with_labels=True, node_color="#99ccff", font_size=8, node_size=2000
+    )
+    plt.show()