'''
Right click to add nodes.
Left drag between node centers to add arrows.
Use up/down arrow keys to change the dampening factor d.

The nodes are labelled by their current PageRank.
Ranks are normalized so that they sum to 1000.
Circle areas are proportional to rankings.
'''

import pygame
import math
import numpy as np

pygame.init()

width, height = 1200, 800
screen = pygame.display.set_mode((width, height))

d=0.85 #dampening parameter

# Colors
background_color = (255, 255, 255)
node_color = (0, 120, 225)
text_color = (255, 255, 255)
edge_color = (100, 100, 100)

node_counter = 0
nodes = []
edges = []
rankings = np.array([])


pygame.display.set_caption("PageRank with d="+str(d))

# Font for node labels
font = pygame.font.Font(None, 24)

def create_adjacency_matrix(nodes,edges,d):
    size = len(nodes)
    matrix = np.zeros((size, size))
    for start, end in edges:
        matrix[end][start] = 1  # Note: start and end are switched because we're filling columns
    
    # Normalize columns to sum to 1, handling dangling nodes
    column_sums = np.sum(matrix, axis=0)
    matrix[:, column_sums == 0] = 1 / size
    matrix /= column_sums + (column_sums == 0)  # Avoid division by zero
    
    # Calculate E, a square matrix with each entry being 1/n
    E = np.ones((size, size)) / size
    
    # Update the matrix according to the formula (1-alpha)*matrix + alpha*E
    matrix = d * matrix + (1-d) * E
    
    return matrix


def power_method(matrix, iterations=100):
    size = matrix.shape[0]
    vector = np.ones(size) / size  # Start with an equal distribution
    for _ in range(iterations):
        vector = np.dot(matrix, vector)
        vector /= np.sum(vector)  # Normalize vector
    return vector

def update_rankings():
    pygame.display.set_caption("PageRank with d="+str(round(d,2)))
    return power_method(create_adjacency_matrix(nodes,edges,d))
    

def draw_nodes_with_rankings():
    for node in nodes:
        radius = int((rankings[node['id']])**0.5 * 150)
        pygame.draw.circle(screen, node_color, node["pos"], radius)
        # Multiply the rank by 1000 and convert to integer
        rank_int = int(rankings[node['id']] * 1000)
        rank_str = f"{rank_int}"  # Convert the integer rank to a string
        text_surface = font.render(rank_str, True, text_color)
        screen.blit(text_surface, (node["pos"][0] - text_surface.get_width() // 2, node["pos"][1] - text_surface.get_height() // 2))




def draw_edges():
    for edge in edges:
        start_pos = nodes[edge[0]]["pos"]
        end_pos = nodes[edge[1]]["pos"]

        # Draw the line from start_pos to end_pos
        pygame.draw.line(screen, edge_color, start_pos, end_pos, 2)

        # Draw arrow head at midpoint of edge
        mid_pos = ((start_pos[0] + end_pos[0]) / 2, (start_pos[1] + end_pos[1]) / 2)
        rotation = math.atan2(end_pos[0] - start_pos[0], end_pos[1] - start_pos[1])
        arrow_size = 15  # Size of the arrowhead
        arrow_angle = math.pi / 6  # Angle at the arrowhead tip
        rotation+=math.pi
        arrow_point2 = (mid_pos[0] + arrow_size * math.sin(rotation + arrow_angle), mid_pos[1] + arrow_size * math.cos(rotation + arrow_angle))
        arrow_point3 = (mid_pos[0] + arrow_size * math.sin(rotation - arrow_angle), mid_pos[1] + arrow_size * math.cos(rotation - arrow_angle))
        pygame.draw.polygon(screen, edge_color, [mid_pos, arrow_point2, arrow_point3])
        
def find_node_at_pos(pos):
    for i, node in enumerate(nodes):
        if (node["pos"][0] - pos[0]) ** 2 + (node["pos"][1] - pos[1]) ** 2 <= 30**2:
            return i
    return None

# Main loop
running = True
is_dragging = False
start_node_index = None

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                rankings=update_rankings()
            elif event.key == pygame.K_UP:
                d+=0.01
                d=min(d,1)
                rankings=update_rankings()
            elif event.key == pygame.K_DOWN:
                d-=0.01
                d=max(d,0)
                rankings=update_rankings()

                
        elif event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 3:  # Right-click to create a node
                nodes.append({"id": node_counter, "pos": event.pos})
                node_counter += 1
                rankings=update_rankings()
            elif event.button == 1:  # Left-click to start drawing an edge
                start_node_index = find_node_at_pos(event.pos)
                if start_node_index is not None:
                    is_dragging = True
        elif event.type == pygame.MOUSEBUTTONUP:
            if event.button == 1 and is_dragging:  # Release left-click to finish the edge
                end_node_index = find_node_at_pos(event.pos)
                if end_node_index is not None and start_node_index != end_node_index:
                    edges.append((start_node_index, end_node_index))
                    rankings=update_rankings()
                is_dragging = False

    # Drawing
    screen.fill(background_color)
    draw_edges()
    draw_nodes_with_rankings()
    pygame.display.flip()

pygame.quit()