"""
Face-area-balanced planar layout for maximal planar graphs.

Starts from a Tutte embedding (outer face on a circle) and uses random-search
optimization to equalize interior face areas while maintaining planarity.
"""
import networkx as nx
import numpy as np
import scipy.linalg


def _get_all_faces(emb):
    """Enumerate all faces of a PlanarEmbedding."""
    seen, faces = set(), []
    for v in emb.nodes():
        for w in emb[v]:
            if (v, w) not in seen:
                face = emb.traverse_face(v, w)
                for i in range(len(face)):
                    seen.add((face[i], face[(i+1) % len(face)]))
                faces.append(tuple(face))
    return faces


def _tutte_layout(G, outer_face):
    """Standard uniform Tutte embedding with outer face on unit circle."""
    outer = list(outer_face)
    k = len(outer)
    fixed_pos = {v: np.array([np.cos(2*np.pi*i/k - np.pi/2),
                              np.sin(2*np.pi*i/k - np.pi/2)])
                 for i, v in enumerate(outer)}
    interior = [v for v in G.nodes() if v not in fixed_pos]
    if not interior:
        return fixed_pos
    int_idx = {v: i for i, v in enumerate(interior)}
    m = len(interior)
    A = np.zeros((m, m)); bx = np.zeros(m); by = np.zeros(m)
    for v in interior:
        i = int_idx[v]; nbrs = list(G.neighbors(v)); deg = len(nbrs)
        A[i, i] = -1.0
        for u in nbrs:
            if u in int_idx:
                A[i, int_idx[u]] += 1.0 / deg
            else:
                bx[i] -= fixed_pos[u][0] / deg
                by[i] -= fixed_pos[u][1] / deg
    px = scipy.linalg.solve(A, bx)
    py = scipy.linalg.solve(A, by)
    pos = dict(fixed_pos)
    for v in interior:
        pos[v] = np.array([px[int_idx[v]], py[int_idx[v]]])
    return pos


def _segments_cross(p1, p2, p3, p4):
    """True iff open segments p1-p2 and p3-p4 cross (not at endpoints)."""
    def cross(o, a, b):
        return (a[0]-o[0])*(b[1]-o[1]) - (a[1]-o[1])*(b[0]-o[0])
    d1 = cross(p3, p4, p1); d2 = cross(p3, p4, p2)
    d3 = cross(p1, p2, p3); d4 = cross(p1, p2, p4)
    return ((d1 > 0 and d2 < 0) or (d1 < 0 and d2 > 0)) and \
           ((d3 > 0 and d4 < 0) or (d3 < 0 and d4 > 0))


def _count_crossings(G, pos):
    """Number of pairs of edges whose straight-line segments cross."""
    edges = list(G.edges())
    n = 0
    for i in range(len(edges)):
        a, b = edges[i]
        for j in range(i+1, len(edges)):
            c, d = edges[j]
            if len({a, b, c, d}) < 4:
                continue
            if _segments_cross(pos[a], pos[b], pos[c], pos[d]):
                n += 1
    return n


def _face_area(pos, face):
    """Unsigned area of a triangular (or polygonal) face."""
    pts = [pos[v] for v in face]
    return 0.5 * abs((pts[1][0]-pts[0][0])*(pts[2][1]-pts[0][1])
                   - (pts[1][1]-pts[0][1])*(pts[2][0]-pts[0][0]))


def _edge_lengths(G, pos):
    return [np.linalg.norm(pos[u] - pos[v]) for u, v in G.edges()]


def _score(G, pos, interior_faces, w_area=1.0, w_edge=0.3):
    """Combined penalty: normalized variance of face areas + edge lengths."""
    areas = np.array([_face_area(pos, f) for f in interior_faces])
    a_mean = areas.mean()
    area_score = np.sum((areas - a_mean)**2) / (a_mean**2)
    lens = np.array(_edge_lengths(G, pos))
    l_mean = lens.mean()
    edge_score = np.sum((lens - l_mean)**2) / (l_mean**2)
    return w_area * area_score + w_edge * edge_score


def balanced_planar_layout(
    G,
    outer_face,
    n_explore=8000,
    n_refine=4000,
    explore_step=0.30,
    refine_step=0.05,
    w_area=1.0,
    w_edge=0.3,
    seed=42,
    verbose=False,
):
    """
    Compute a planar layout for a maximal planar graph G whose interior faces
    have roughly equal area and whose edges have roughly equal length.

    Starts from a uniform Tutte embedding (outer face on a unit circle in CCW
    order starting at the south pole), then runs random-search optimization,
    accepting only moves that keep the layout planar.

    Parameters
    ----------
    G : networkx.Graph
        A simple connected planar graph; intended for maximal planar graphs.
    outer_face : tuple
        Cyclic ordering of vertices on the outer face. Pinned during
        optimization.
    n_explore, n_refine : int
        Number of iterations in the exploration and refinement phases.
    explore_step, refine_step : float
        Standard deviation of Gaussian moves in each phase.
    w_area, w_edge : float
        Weights for the face-area and edge-length penalties in the score.
    seed : int
        RNG seed for reproducibility.
    verbose : bool
        If True, print progress.

    Returns
    -------
    dict[node, np.ndarray]
        Position of each vertex.
    """
    is_planar, emb = nx.check_planarity(G)
    if not is_planar:
        raise ValueError("G is not planar")

    pos = _tutte_layout(G, outer_face)

    all_faces = _get_all_faces(emb)
    outer_set = set(outer_face)
    interior_faces = [f for f in all_faces if set(f) != outer_set]

    if _count_crossings(G, pos) != 0:
        raise RuntimeError("Initial Tutte layout is not planar; check inputs")

    interior_nodes = [v for v in G.nodes() if v not in outer_set]
    if not interior_nodes:
        return pos

    rng = np.random.default_rng(seed)
    best_score = _score(G, pos, interior_faces, w_area, w_edge)
    best_pos = {k: v.copy() for k, v in pos.items()}
    initial_score = best_score

    for phase_name, n_iter, step in [("explore", n_explore, explore_step),
                                     ("refine",  n_refine,  refine_step)]:
        pos = {k: v.copy() for k, v in best_pos.items()}
        accepted = 0
        for it in range(n_iter):
            v = interior_nodes[rng.integers(len(interior_nodes))]
            delta = rng.normal(0, step, 2)
            old = pos[v].copy()
            pos[v] = old + delta
            if _count_crossings(G, pos) == 0:
                new_score = _score(G, pos, interior_faces, w_area, w_edge)
                if new_score < best_score:
                    best_score = new_score
                    best_pos = {k: vv.copy() for k, vv in pos.items()}
                    accepted += 1
                else:
                    pos[v] = old
            else:
                pos[v] = old
        if verbose:
            print(f"  {phase_name}: score {best_score:.4f}, "
                  f"accepted {accepted}/{n_iter}")

    if verbose:
        print(f"Initial score: {initial_score:.4f}")
        print(f"Final score:   {best_score:.4f}")

    return best_pos


# ── Demo ─────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    edges = [(0,1),(0,2),(0,3),(0,4),(0,5),(0,6),
             (1,2),(1,3),(1,4),(1,5),(1,6),
             (2,3),(2,4),(3,5),(4,6)]
    G = nx.Graph()
    G.add_nodes_from(range(7))
    G.add_edges_from(edges)

    pos = balanced_planar_layout(G, outer_face=(0, 3, 5), verbose=True)
    print("\nFinal positions:")
    for v in sorted(pos):
        p = pos[v]
        print(f"  {v}: ({p[0]:+.3f}, {p[1]:+.3f})")
    print(f"Crossings: {_count_crossings(G, pos)}")
