genesis/visualization.py

# genesis/visualization.py
import os
from neo4j import GraphDatabase
import plotly.graph_objects as go

# =========================
# CONFIGURATION
# =========================
NEO4J_URI = os.getenv("NEO4J_URI")
NEO4J_USER = os.getenv("NEO4J_USER")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")

COLOR_PATHWAY_NODE = "#00FF99"  # Biotech green
COLOR_FUNDING_NODE = "#FF9900"  # Orange
COLOR_EDGE = "#AAAAAA"

# =========================
# NEO4J CONNECTION
# =========================
driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))

def run_neo4j_query(query, params=None):
    """Run a Cypher query and return results."""
    with driver.session() as session:
        return list(session.run(query, params or {}))

# =========================
# GRAPH UTILS
# =========================
def create_plotly_graph(nodes, edges, title):
    """Creates an interactive Plotly network graph."""
    node_x, node_y, node_text, node_color = [], [], [], []
    edge_x, edge_y = [], []

    # Layout algorithm (spring layout for better spacing)
    import networkx as nx
    G = nx.Graph()
    for node_id, label, color in nodes:
        G.add_node(node_id, label=label, color=color)
    for src, dst in edges:
        G.add_edge(src, dst)

    pos = nx.spring_layout(G, seed=42, k=0.5)

    for src, dst in edges:
        x0, y0 = pos[src]
        x1, y1 = pos[dst]
        edge_x.extend([x0, x1, None])
        edge_y.extend([y0, y1, None])

    for node_id, label, color in nodes:
        x, y = pos[node_id]
        node_x.append(x)
        node_y.append(y)
        node_text.append(label)
        node_color.append(color)

    edge_trace = go.Scatter(
        x=edge_x, y=edge_y,
        line=dict(width=1, color=COLOR_EDGE),
        hoverinfo='none',
        mode='lines'
    )

    node_trace = go.Scatter(
        x=node_x, y=node_y,
        mode='markers+text',
        text=node_text,
        textposition="top center",
        hoverinfo="text",
        marker=dict(
            color=node_color,
            size=18,
            line=dict(width=2, color='#000000')
        )
    )

    fig = go.Figure(data=[edge_trace, node_trace])
    fig.update_layout(
        title=title,
        title_x=0.5,
        plot_bgcolor="#111111",
        paper_bgcolor="#111111",
        font=dict(color="#FFFFFF"),
        showlegend=False,
        margin=dict(l=10, r=10, t=40, b=10)
    )
    return fig

# =========================
# PATHWAY GRAPH
# =========================
def generate_pathway_graph(pathway_name):
    """
    Generate an interactive graph for a given metabolic pathway.
    """
    query = """
    MATCH (p:Pathway {name: $name})-[r:INVOLVES]->(m:Molecule)
    RETURN p.name AS pathway, m.name AS molecule
    """
    results = run_neo4j_query(query, {"name": pathway_name})

    if not results:
        return None

    nodes = [(pathway_name, pathway_name, COLOR_PATHWAY_NODE)]
    edges = []
    for record in results:
        mol_name = record["molecule"]
        nodes.append((mol_name, mol_name, "#00BFFF"))  # Blue for molecules
        edges.append((pathway_name, mol_name))

    return create_plotly_graph(nodes, edges, f"Metabolic Pathway: {pathway_name}")

# =========================
# FUNDING NETWORK
# =========================
def generate_funding_network(industry_keyword):
    """
    Generate an interactive funding network graph for companies in a given biotech domain.
    """
    query = """
    MATCH (c:Company)-[f:FUNDED_BY]->(i:Investor)
    WHERE toLower(c.industry) CONTAINS toLower($keyword)
    RETURN c.name AS company, i.name AS investor
    """
    results = run_neo4j_query(query, {"keyword": industry_keyword})

    if not results:
        return None

    nodes = []
    edges = []
    seen_nodes = set()

    for record in results:
        comp = record["company"]
        inv = record["investor"]

        if comp not in seen_nodes:
            nodes.append((comp, comp, COLOR_FUNDING_NODE))
            seen_nodes.add(comp)
        if inv not in seen_nodes:
            nodes.append((inv, inv, "#FFD700"))  # Gold for investors
            seen_nodes.add(inv)

        edges.append((comp, inv))

    return create_plotly_graph(nodes, edges, f"Funding Network: {industry_keyword}")

# =========================
# CLEANUP
# =========================
def close_driver():
    driver.close()