import math

import matplotlib.pyplot as plt
import numpy as np
from graph_tool.all import *

from src import centrality
from src import plot
from src import fitting

def random_graph(n=5000, seed=None):
    """
    Uniformly random point cloud generation.
    `n` [int] Number of points to generate. Default 5000 seems like a good starting point in point density and corresponding runtime for the subsequent calculations.
    @return [numpy.ndarray] Array of shape(n, 2) containing the coordinates for each point of the generated point cloud.
    """
    if seed is None:
        import secrets
        seed = secrets.randbits(128)
    rng = np.random.default_rng(seed=seed)
    return rng.random((n, 2)), seed


def spatial_graph(adata):
    g, pos = graph_tool.generation.triangulation(adata, type="delaunay")
    g.vp["pos"] = pos
    weight = g.new_edge_property("double")
    for e in g.edges():
        weight[e] = math.sqrt(sum(map(abs, pos[e.source()].a - pos[e.target()].a)))**2
    return g, weight


points, seed = random_graph()
g, weight = spatial_graph(points)
g = GraphView(g)

# calculate centrality values
vp = closeness(g, weight=weight)
vp.a = np.nan_to_num(vp.a) # correct floating point values
# ep.a = np.nan_to_num(ep.a) # correct floating point values

# calculate convex hull
convex_hull = centrality.convex_hull(g)

# plot graph with convex_hull
fig = plt.figure(figsize=(15, 5))
ax0, ax1 = fig.subplots(1, 2)
plot.graph_plot(fig, ax0, g, vp, convex_hull, f"Closeness without prediction")

# generate model based on convex hull and associated centrality values
quantification = plot.quantification_data(g, vp, convex_hull)

# optimize model's piece-wise linear function
d = quantification[:, 0]
C = quantification[:, 1]
m_opt, c0_opt, b_opt, aic_opt = fitting.fit_piece_wise_linear(d, C)

vp = centrality.correct(g, vp, m_opt, c0_opt, b_opt)
plot.graph_plot(fig, ax1, g, vp, convex_hull, f"Closeness with model prediction")

fig.savefig(f"model_prediction_comparison.svg", format='svg')