"""Small complex cross-ratio lattice utilities."""
import numpy as np


def complete_d(a, b, c, eps=1e-12):
    """Return d such that (a,b,c,d) has cross-ratio -1.

    The returned singular mask identifies near-zero denominators; singular
    values are represented by complex infinity rather than NaN.
    """
    den = (b - c) - (a - b)
    num = (b - c) * a - (a - b) * c
    singular = np.abs(den) < eps
    d = np.empty(np.broadcast(a, b, c).shape, dtype=np.result_type(a, b, c))
    np.divide(num, den, out=d, where=~singular)
    d = np.where(singular, np.inf + 0j, d)
    return d, singular


def complete_c(a, b, d, eps=1e-12):
    """Inverse completion: solve the same equation for c."""
    # From d=((b-c)a-(a-b)c)/((b-c)-(a-b)), collect c.
    # c = ((a-b)*d - (b)*a + b*a? use direct Möbius inversion below)
    # Cross-ratio equation gives (a-b)(c-d)+(b-c)(d-a)=0.
    # c*((a-b)-(d-a)) = (a-b)d - b*(d-a)
    den = (a - b) - (d - a)
    num = (a - b) * d - b * (d - a)
    singular = np.abs(den) < eps
    c = np.empty(np.broadcast(a, b, d).shape, dtype=np.result_type(a, b, d))
    np.divide(num, den, out=c, where=~singular)
    c = np.where(singular, np.inf + 0j, c)
    return c, singular


def cross_ratio(a, b, c, d):
    return (a-b)*(c-d)/((b-c)*(d-a))


def residual(a, b, c, d, eps=1e-14):
    den = (b-c)*(d-a)
    out = np.full(np.broadcast(a,b,c,d).shape, np.inf, dtype=float)
    good = np.abs(den) > eps
    out[good] = np.abs(cross_ratio(a,b,c,d)[good] + 1)
    return out


def generate_lattice(n, seed=0, dtype=np.complex128):
    """Generate a lattice from random first row/column and exact completion."""
    rng = np.random.default_rng(seed)
    z = np.zeros((n,n), dtype=dtype)
    z[0, :] = (rng.uniform(-1,1,n) + 1j*rng.uniform(-1,1,n)).astype(dtype)
    z[:, 0] = (rng.uniform(-1,1,n) + 1j*rng.uniform(-1,1,n)).astype(dtype)
    # Avoid an exactly shared random origin while retaining a valid boundary.
    z[0,0] = 0.2 + 0.15j
    singular = 0
    for i in range(n-1):
        for j in range(n-1):
            c, bad = complete_c(z[i,j], z[i+1,j], z[i,j+1])
            if bad: singular += 1
            z[i+1,j+1] = c
    return z, singular


def plaquette_residuals(z):
    return residual(z[:-1,:-1], z[1:,:-1], z[1:,1:], z[:-1,1:])


def constrained_from_boundary(top, left):
    n = len(top)
    z = np.zeros((n,n), dtype=np.result_type(top, left))
    z[0,:], z[:,0] = top, left
    for i in range(n-1):
        for j in range(n-1):
            z[i+1,j+1], _ = complete_c(z[i,j], z[i+1,j], z[i,j+1])
    return z


def additive_from_boundary(top, left):
    n = len(top)
    z = np.zeros((n,n), dtype=np.result_type(top, left))
    z[0,:], z[:,0] = top, left
    for i in range(n-1):
        for j in range(n-1):
            z[i+1,j+1] = z[i+1,j] + z[i,j+1] - z[i,j]
    return z