# pyright: reportIndexIssue=false, reportMissingParameterType=false, reportUnknownLambdaType=false

import hashlib
import sys

import networkx as nx
import numpy as np
import rerun as rr
import yaml


def start():
    # Give relative path or default to local pixi.lock
    lockfile_path = sys.argv[1] if len(sys.argv) > 1 else "pixi.lock"

    with open(lockfile_path) as file:
        lockfile_data = yaml.safe_load(file)

    package_data = lockfile_data["packages"]

    graph = nx.DiGraph()
    for package in package_data:
        if "linux" in package.get("conda", ""):
            package_url = package.get("conda") or package.get("pypi")
            package_name = (
                package_url.split("/")[-1].split("-")[0] if package_url else "unknown_package"
            )
            dependencies = package.get("depends", [])
            graph.add_node(package_name)
            for dep in dependencies:
                graph.add_edge(package_name, dep.split(" ")[0])

    rr.init("fdg", spawn=True)

    # Identify the node with the highest degree
    central_node = max(graph.degree, key=lambda x: x[1])[0]

    # Initial positions with the central node at the center
    initial_pos = nx.spring_layout(graph, dim=3)
    initial_pos[central_node] = np.array([0.5, 0.5, 0.5])  # Center position

    # Apply the force-directed simulation
    apply_forces_and_log(graph, initial_pos)


def hash_string_to_int(string):
    return int(hashlib.sha256(string.encode("utf-8")).hexdigest(), 16) % (10**8)


# Function to get color
def get_color_for_node(node, color_cache):
    if node not in color_cache:
        np.random.seed(hash_string_to_int(node))
        color_cache[node] = np.random.rand(3)  # Generate and store color
    return color_cache[node]


def apply_forces_and_log(graph, pos):
    damping = 0.9
    max_force = 1
    degree_scale = 2  # Scale factor for degree-based forces
    dist_scale = 0.5

    iterations = 1000
    repulsive_force = 0.01
    attractive_force = 0.005

    # Memoization dictionary
    color_cache = {}

    for iteration in range(iterations):
        force = {node: np.zeros(3) for node in graph}

        # Degree-based repulsive forces
        for i, node1 in enumerate(graph):
            for node2 in list(graph)[i + 1 :]:
                diff = pos[node1] - pos[node2]
                dist = (np.linalg.norm(diff) + 1e-9) * dist_scale
                degree_factor = (graph.degree(node1) + graph.degree(node2)) * degree_scale
                repel = repulsive_force * degree_factor / dist**2
                force_vector = repel * diff  # / dist
                force[node1] += np.clip(force_vector, -max_force, max_force)
                force[node2] -= np.clip(force_vector, -max_force, max_force)

        # Degree-based attractive forces
        for edge in graph.edges():
            u, v = edge
            diff = pos[u] - pos[v]
            dist = dist = (np.linalg.norm(diff) + 1e-9) * dist_scale
            if dist > 0:
                degree_factor = (graph.degree(u) + graph.degree(v)) * degree_scale
                attract = (attractive_force * dist**2) / degree_factor
                force[u] -= attract * diff / dist
                force[v] += attract * diff / dist

        # Update positions with damping
        if iteration % 1 == 0:
            for node in graph:
                pos[node] += force[node] * damping
                position = np.array(pos[node])
                color = get_color_for_node(node, color_cache)  # Retrieve color, memoized
                rr.log(
                    f"graph_nodes/{node}",
                    rr.Points3D(
                        [position], colors=[color], radii=max(graph.degree(node) / 20, 1.5)
                    ),
                    rr.AnyValues(node),
                )

            edges_array = np.array([[pos[u], pos[v]] for u, v in graph.edges()])

            # Log the edges array
            rr.log(
                "graph_nodes/graph_edges",
                rr.LineStrips3D(edges_array, radii=0.2, colors=[1, 1, 1, 0.1]),
            )

    return pos


if __name__ == "__main__":
    try:
        start()
    except KeyboardInterrupt:
        print("\nInterrupted by user, exiting...")
        sys.exit(0)
