import psycopg2
import random
import numpy as np
import matplotlib.pyplot as plt
import networkx as nx

# PostgreSQL ühenduse andmed
db_params = {
    'host': 'localhost',
    'port': '5432',
    'dbname': 'data',
    'user': 'osm',
    'password': 'osm'
}

# Ühendu andmebaasiga
conn = psycopg2.connect(**db_params)
cur = conn.cursor()

# Lae graafi servad
cur.execute("""
    SELECT id, source, target, cost FROM aaa_osm.osm_roads
    WHERE source IS NOT NULL AND target IS NOT NULL
""")
edges = cur.fetchall()

# Ehita NetworkX graaf
G = nx.DiGraph()
for edge_id, source, target, cost in edges:
    G.add_edge(source, target, weight=cost)

# Valime juhuslikud lähte- ja sihtpunktid
nodes = list(G.nodes)
num_pairs = 100
random.seed(42)
origins = random.sample(nodes, num_pairs)
destinations = random.sample(nodes, num_pairs)

# Määrame igale punktile "mass" (nt rahvaarv)
P = {node: random.randint(100, 1000) for node in nodes}

# Arvutame gravitatsioonimudeli T_ij = (P_i * P_j) / D_ij^2
gravity_matrix = np.zeros((num_pairs, num_pairs))
for i, origin in enumerate(origins):
    for j, dest in enumerate(destinations):
        if origin == dest:
            continue
        try:
            length = nx.shortest_path_length(G, source=origin, target=dest, weight='weight')
            gravity_matrix[i, j] = (P[origin] * P[dest]) / (length ** 2)
        except nx.NetworkXNoPath:
            gravity_matrix[i, j] = 0

# Visualiseerime tulemuse
plt.figure(figsize=(10, 8))
plt.imshow(gravity_matrix, cmap='hot', interpolation='nearest')
plt.colorbar(label='Gravitatsiooniline intensiivsus (T_ij)')
plt.title('Gravitatsioonimudel (100 lähtekohta ja sihtkohta)')
plt.xlabel('Sihtkoht')
plt.ylabel('Lähtekoht')
plt.tight_layout()
plt.savefig("gravity_model_heatmap.png")
plt.show()

# Sulgeme ühenduse
cur.close()
conn.close()