"""Finite-grid polarization-gap and spherical-triangle Chern monitor."""
import numpy as np


def spherical_area(a, b, c):
    num = np.einsum('...i,...i->...', a, np.cross(b, c))
    den = 1 + np.einsum('...i,...i->...', a, b) + np.einsum('...i,...i->...', b, c) + np.einsum('...i,...i->...', c, a)
    return 2 * np.arctan2(num, den)


def monitor_field(m, return_triangles=False):
    """Return (minimum norm, Chern estimate) for periodic KxK vector field."""
    m = np.asarray(m, dtype=float)
    norms = np.linalg.norm(m, axis=-1)
    gap = float(norms.min())
    if gap <= 1e-14:
        return (gap, np.nan) if not return_triangles else (gap, np.nan, [])
    n = m / norms[..., None]
    K = m.shape[0]
    total, areas = 0.0, []
    for i in range(K):
        for j in range(K):
            a, b = n[i,j], n[(i+1)%K,j]
            c, d = n[(i+1)%K,(j+1)%K], n[i,(j+1)%K]
            for tri in ((a,b,c), (a,c,d)):
                ar = float(spherical_area(*tri)); total += ar; areas.append(ar)
    result = (gap, total/(4*np.pi))
    return result + (areas,) if return_triangles else result


def qwz_field(K, mass, anisotropy=1.0):
    """Periodic response field with gap closings at mass=-2, 0, 2."""
    x = 2*np.pi*np.arange(K)/K
    a, t = np.meshgrid(x, x, indexing='ij')
    return np.stack([np.sin(a), anisotropy*np.sin(t), mass + np.cos(a) + np.cos(t)], axis=-1)


def fixed_point_parity(mass):
    # z signs at (0,0),(pi,0),(0,pi),(pi,pi), with nonzero perturbation avoided.
    zs = [mass+2, mass, mass, mass-2]
    return int(np.sign(np.prod(zs)))