#!/usr/bin/env python
# coding: utf-8

# adapted originally from https://gist.github.com/ericdill/9942ac55c2c9f6a550973dd2dc3653a4
# and then later from https://gist.github.com/msarahan/f425082c67bd985cb4edd0edb4d50de0
# 
# updated for repodata format changes

# In[1]:


import copy
import json
import glob
from itertools import chain
from pathlib import Path

from matplotlib import pyplot as plt
import networkx as nx


# In[2]:


import requests_cache
requests_cache.install_cache()
import requests

channels = [
    "conda-forge/linux-64",
    "conda-forge/noarch",
]

dg = nx.DiGraph()

for channel in channels:
    url = f"https://conda.anaconda.org/{channel}/repodata.json"
    r = requests.get(url)
    assert r.status_code == 200, r.status_code
    repodata = json.loads(r.text)
    for meta in chain(repodata['packages'].values(), repodata['packages.conda'].values()):
        dg.add_node(meta['name'])
        for dep in meta.get('depends', []):
            dg.add_edge(meta['name'], dep.partition(' ')[0])    


# In[3]:


dg.number_of_nodes()


# In[4]:


[name for name in dg if ('dolfin' in name or 'fenics' in name)]


# In[5]:


def dependents(G: nx.DiGraph, roots: str | list[str], steps: int=10) -> nx.DiGraph:
    """compute dependent packages, given one or more root packages"""
    if isinstance(roots, str):
        roots = [roots]
    neighbors = set(roots)
    for n in range(steps):
        immediate_neighbors = copy.copy(neighbors)
        for neighbor in immediate_neighbors:
            neighbors.update(G.predecessors(neighbor))
    return nx.subgraph(G, neighbors)


# In[6]:


sg = dependents(dg, ['fenics-ufl', 'fenics-ufcx'])


# In[7]:


def topo_pos(G: nx.DiGraph) -> dict:
    """Display in topological order, with simple offsetting for legibility"""
    pos_dict = {}
    for i, node_list in enumerate(nx.topological_generations(G)):
        x_offset = len(node_list) / 2
        y_offset = 0.5 / len(node_list)
        for j, name in enumerate(node_list):
            pos_dict[name] = (j - x_offset, -i)

    return pos_dict

sgr = sg.reverse()
nx.draw(sgr, topo_pos(sgr), with_labels=True, node_color='#A0CBE2', edge_color='#FFCBE2')


# In[8]:


nx.write_network_text(sgr)