# ezga.utils.symmetry_operations_generator.py
import numpy as np

# ---- Primitive transforms -------------------------------------------------
def _unit(v):
    v = np.asarray(v, float)
    n = np.linalg.norm(v)
    if n == 0: raise ValueError("Zero axis/normal")
    return v / n

def rot_axis_angle(axis, theta):
    """Proper rotation R ∈ SO(3) about `axis` by angle `theta`."""
    u = _unit(axis)
    ux, uy, uz = u
    c, s = np.cos(theta), np.sin(theta)
    C = 1 - c
    return np.array([
        [c+ux*ux*C,   ux*uy*C - uz*s, ux*uz*C + uy*s],
        [uy*ux*C + uz*s, c+uy*uy*C,   uy*uz*C - ux*s],
        [uz*ux*C - uy*s, uz*uy*C + ux*s, c+uz*uz*C  ]
    ], float)

def mirror(normal):
    """Reflection across plane with normal vector `normal` (det = -1)."""
    n = _unit(normal)
    # Householder reflection: H = I - 2 n n^T
    return np.eye(3) - 2.0*np.outer(n, n)

def inversion():
    """Spatial inversion (parity)."""
    return -np.eye(3)

def improper(axis, n):
    """S_n: rotation by 2π/n about axis followed by mirror ⟂ axis."""
    R = rot_axis_angle(axis, 2*np.pi/n)
    M = mirror(axis)  # plane perpendicular to axis
    return M @ R

# ---- Group closure from generators ---------------------------------------
def group_closure(generators, tol=1e-10, max_elems=512):
    """
    Return the finite group generated by `generators` under matrix multiplication.
    Deduplicates by Frobenius norm within `tol`.
    """
    def _canon(M):
        # enforce exact orthogonality numerically for stability
        U, _, Vt = np.linalg.svd(M)
        M = U @ Vt
        return M

    ops = []
    def _contains(M):
        return any(np.linalg.norm(M - N, ord='fro') < tol for N in ops)

    # start with identity
    I = np.eye(3)
    ops.append(I)
    frontier = [I]

    gens = [ _canon(G) for G in generators ]

    while frontier and len(ops) < max_elems:
        new_frontier = []
        for A in frontier:
            for G in gens:
                M = _canon(G @ A)
                if not _contains(M):
                    ops.append(M)
                    new_frontier.append(M)
        frontier = new_frontier
    return ops

# ---- Prebuilt helpers for common groups ----------------------------------
def build_Cn(n, axis=(0,0,1), include_inversion=False, include_mirror_axial=False, include_mirror_horizontal=False):
    gens = [rot_axis_angle(axis, 2*np.pi/n)]
    if include_inversion:
        gens.append(inversion())
    if include_mirror_axial:      # adds Cnv / Dnh families with extra mirrors
        # pick one vertical mirror plane that contains axis and x
        gens.append(mirror(np.cross(axis, (0,1,0)) or (1,0,0)))
    if include_mirror_horizontal:
        gens.append(mirror(axis))
    return group_closure(gens)

def build_Dn(n, axis=(0,0,1), include_inversion=False, include_horizontal_mirror=False):
    # Dn generated by: Cn about z and C2 about x (⊥ axis).
    gens = [
        rot_axis_angle(axis, 2*np.pi/n),
        rot_axis_angle((1,0,0), np.pi)
    ]
    if include_inversion:
        gens.append(inversion())
    if include_horizontal_mirror:
        gens.append(mirror(axis))
    return group_closure(gens)

def build_Oh():
    # Octahedral group with inversion (O_h, order 48).
    # Generators: 90° about x; 120° about (1,1,1); inversion.
    gens = [
        rot_axis_angle((1,0,0), np.pi/2),
        rot_axis_angle((1,1,1), 2*np.pi/3),
        inversion()
    ]
    return group_closure(gens, max_elems=64)

def build_Ih():
    # Icosahedral full group (I_h, order 120) – expensive but available.
    # Practical generator set (e.g., 5-fold about z and 3-fold about (1,1,1), plus inversion).
    phi = (1 + np.sqrt(5))/2
    gens = [
        rot_axis_angle((0,0,1), 2*np.pi/5),
        rot_axis_angle((1,1,1), 2*np.pi/3),
        inversion()
    ]
    return group_closure(gens, max_elems=128)
