# venture_networks_gradio.py

import networkx as nx
import matplotlib.pyplot as plt
import gradio as gr
from io import BytesIO

# Define investors and their companies
investors = {
    "Accel": ["Meta", "Dropbox", "Spotify", "Adroll", "PackLink", "Zoom", "Slack"],
    "Andreessen Horowitz": [
        "Airbnb", "Lyft", "Pinterest", "Accolade", "Actifio", "Affirm", "Amify",
        "Apptio", "Arimo", "Asana", "Asserts", "Authy", "Automatic", "Barefoot Networks",
        "Bebop", "BeReal", "Blockspring", "Boku", "Box", "Bracket Computing", "Bump",
        "BuzzFeed", "Capital", "Cardiogram", "Caviar", "Citiizen", "CipherCloud",
        "ClearStory", "Coinbase", "Cumulus", "DeepMap", "Delphix", "DigitalOcean",
        "Dollar Shave Club", "Earn", "EQRx", "Erasca", "Meta", "FiftyThree", "Flutter",
        "Fusion-io", "Gainspeed", "Genius", "GitHub", "Gobble", "Golden", "Greenlight",
        "Granular", "Greplify", "Groupon", "Hatch", "iCracked", "Imgur", "Instacart",
        "Instagram", "Intrinsic", "Isovalent", "Jopwell", "Julep", "Jungla", "Keybase",
        "Kno", "Knowable", "Knowmia", "Loom", "Lowkey", "Luma", "Lyrebird", "Lytro",
        "Massive", "Morta Security", "Nautilus Biotechnology", "Nicira", "Oculus", "Okta",
        "Onshape", "Opendoor", "OpenInvest", "Optimizely", "Optimyze", "PagerDuty",
        "Parsec", "Patch Biosciences", "PatientPing", "PicCloud", "Pinterest", "Pixate",
        "Platfora", "Product Hunt", "Quantopian", "RHL Ventures", "Reddit", "Reflektive",
        "Rigetti", "Rival", "Robinhood", "Roblox", "Rockmelt", "RTFKT", "Samsara", "Seed",
        "Shapeways", "ShoeDazzle", "SignalFx", "SigOpt", "Silver Tail Systems", "Skout",
        "Skype", "Slack", "Socialcam", "Soylent", "Stack Overflow", "Standard Treasury",
        "Swell", "Swiftype", "Tabular", "Talaria", "Tasty Labs", "Teleport", "Tenfold",
        "Tidemark", "Tilt", "TinyCo", "Tmunity", "Tomfoolery", "Wise", "TXN", "Urban Engines",
        "Usermind", "Viki", "Walker and Company", "Wit.ai", "Within", "Zynga", "Altitude Learning",
        "Bromium", "Factual", "Instart", "Proven", "Rabbit", "UnifyID"
    ],
    "Google Ventures": ["Uber", "LendingClub"],
    "Greylock": ["Workday", "Palo Alto Networks"],
    "Lightspeed Venture Partners": ["Snap", "Grubhub", "AppDynamics"],
    "Benchmark": ["Snap", "Uber", "WeWork"],
    "Norwest Venture Partners": ["LendingClub", "Opendoor"],
    "Emergence Capital Partners": ["Zoom", "Box", "Salesforce"],
    "Trinity Ventures": ["New Relic", "Care.com", "TubeMogul"],
    "Citi Ventures": ["Square", "Nutanix"],
    "Sequoia": [
        "Alphabet (Google)", "NVIDIA", "Dropbox", "Airbnb", "Affirm",
        "Airbnb", "Eventbrite", "Meta", "Lyft", "Opendoor", "Palantir", "Reddit", "Uber",
        "23ANDME", "AMPLITUDE", "APPLE", "ATARI", "AURORA", "BARRACUDA", "BIRD", "BLOCK",
        "BRIDGEBIO", "CISCO", "CONFLUENT", "DOORDASH", "ELECTRONIC ARTS", "FIREEYE", "GOOGLE",
        "GREEN DOT", "GUARDANT HEALTH", "HUBSPOT", "INSTACART", "LINKEDIN", "MEDALLIA", "MONGODB",
        "NATERA", "NETAPP", "NIMBLE STORAGE", "NUBANK", "OKTA", "ORACLE", "PALO ALTO NETWORKS",
        "PAYPAL", "QUALTRICS", "QUANTENNA", "RINGCENTRAL", "ROBINHOOD", "RUCKUS NETWORKS",
        "SERVICENOW", "SNOWFLAKE", "SUMO LOGIC", "SUNRUN", "TRULIA", "UIPATH", "UNITY", "YAHOO"
    ],
    "Y Combinator": ["Dropbox", "Airbnb", "Coinbase", "DoorDash", "Reddit", "Ginkgo Bioworks", "GitLab", "Instacart"]
}

# Define a color map for the investors
investor_colors = {
    "Accel": "blue",
    "Andreessen Horowitz": "green",
    "Google Ventures": "red",
    "Greylock": "purple",
    "Lightspeed Venture Partners": "orange",
    "Benchmark": "brown",
    "Norwest Venture Partners": "pink",
    "Emergence Capital Partners": "yellow",
    "Trinity Ventures": "cyan",
    "Citi Ventures": "magenta",
    "Sequoia": "lime",
    "Y Combinator": "black"
}

# Function to generate the graph image
def generate_graph(selected_investors):
    if not selected_investors:
        selected_investors = list(investors.keys())

    G = nx.Graph()

    # Add edges based on selected investors
    for investor in selected_investors:
        companies = investors[investor]
        for company in companies:
            G.add_edge(investor, company, color=investor_colors[investor])

    # Get edge colors
    edge_colors = [G[u][v]['color'] for u, v in G.edges]

    # Set node colors: investors have their specific color, companies have a default color
    node_colors = []
    for node in G.nodes:
        if node in investor_colors:
            node_colors.append(investor_colors[node])
        else:
            node_colors.append("#F0E68C")  # Khaki for companies

    # Create plot
    plt.figure(figsize=(18, 18))
    pos = nx.spring_layout(G, k=0.2, seed=42)  # Fixed seed for consistency
    nx.draw(
        G, pos,
        with_labels=True,
        node_size=3000,
        node_color=node_colors,
        font_size=7,
        font_weight="bold",
        edge_color=edge_colors,
        width=2
    )
    plt.title("Venture Funded Companies as a Densely Connected Subgraph", fontsize=20)
    plt.axis('off')

    # Save plot to a BytesIO object
    buf = BytesIO()
    plt.savefig(buf, format="png", bbox_inches="tight")
    plt.close()
    buf.seek(0)
    return buf

# Define Gradio interface
def main():
    # Create a sorted list of investors for better UX
    investor_list = sorted(investors.keys())

    iface = gr.Interface(
        fn=generate_graph,
        inputs=gr.CheckboxGroup(
            choices=investor_list,
            label="Select Investors",
            value=investor_list  # Default to all selected
        ),
        outputs=gr.Image(type="pil", label="Venture Network Graph"),
        title="Venture Networks Visualization",
        description="Select investors to visualize their investments in various companies. The graph shows connections between investors and the companies they've invested in.",
        flagging_mode="never"
    )

    iface.launch()

if __name__ == "__main__":
    main()