# 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.0 - 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

# ---- Small vector helpers -------------------------------------------------
def _any_perp(u):
    """Return a unit vector perpendicular to unit vector u."""
    u = _unit(u)
    # pick a reference not collinear with u
    if abs(u[0]) < 0.9:
        ref = np.array([1.0, 0.0, 0.0])
    else:
        ref = np.array([0.0, 1.0, 0.0])
    v = np.cross(u, ref)
    return _unit(v)

def _vertical_mirror_normal(axis, angle=0.0):
    """
    Normal of a vertical mirror plane (containing `axis`).
    The plane is defined by `axis` and a perpendicular direction rotated by `angle` about `axis`.
    """
    u = _unit(axis)
    v = _any_perp(u)                          # some vector ⟂ axis in the plane
    R = rot_axis_angle(u, angle)
    v_rot = R @ v
    n = np.cross(u, v_rot)                    # normal ⟂ to both axis and v_rot
    return _unit(n)

# ---- 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`. Orthogonalizes each element.
    """
    def _canon(M):
        # enforce exact orthogonality numerically for stability
        U, _, Vt = np.linalg.svd(M)
        M = U @ Vt
        # Ensure det is close to ±1 (flip if needed due to SVD ambiguity)
        if np.linalg.det(M) < -1 + 1e-12:
            # reflect last column to fix determinant sign if wildly off; keep orthogonality
            U[:, -1] *= -1
            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

# ---- Core cyclic/dihedral families ---------------------------------------
# ---- Prebuilt helpers for common groups ----------------------------------
def build_Cn(n, axis=(0,0,1), include_inversion=False,
             include_mirror_axial=False, include_mirror_horizontal=False):
    """
    Base Cn with optional inversion (Ci), vertical mirrors (Cnv), horizontal mirror (Cnh).
    Use dedicated builders below for clarity.
    """
    axis = _unit(axis)
    gens = [rot_axis_angle(axis, 2*np.pi/n)]
    if include_inversion:
        gens.append(inversion())
    if include_mirror_axial:
        # add ONE vertical mirror; closure will generate the n equivalent planes
        nrm = _vertical_mirror_normal(axis, angle=0.0)
        gens.append(mirror(nrm))
    if include_mirror_horizontal:
        gens.append(mirror(axis))
    return group_closure(gens, max_elems=4*max(2, n))

def build_Dn(n, axis=(0,0,1), include_inversion=False, include_horizontal_mirror=False):
    """
    Dn generated by: Cn about `axis` and a C2 about an axis ⟂ `axis`.
    Optionally add inversion or σ_h to reach Dnh; use dedicated Dnh/Dnd helpers for clarity.
    """
    axis = _unit(axis)
    # pick one C2 axis perpendicular to `axis`
    perp = _any_perp(axis)
    gens = [
        rot_axis_angle(axis, 2*np.pi/n),
        rot_axis_angle(perp, np.pi)
    ]
    if include_inversion:
        gens.append(inversion())
    if include_horizontal_mirror:
        gens.append(mirror(axis))
    return group_closure(gens, max_elems=8*max(2, n))

# ---- Convenience wrappers for standard notations --------------------------
def build_Cnv(n, axis=(0,0,1)):
    """Cnv (order 2n)."""
    return build_Cn(n, axis=axis, include_mirror_axial=True)

def build_Cnh(n, axis=(0,0,1)):
    """Cnh (order 2n)."""
    return build_Cn(n, axis=axis, include_mirror_horizontal=True)

def build_Dnh(n, axis=(0,0,1)):
    """Dnh (order 4n)."""
    return build_Dn(n, axis=axis, include_horizontal_mirror=True)

def build_Dnd(n, axis=(0,0,1)):
    """
    Dnd (order 4n): Dn plus *dihedral* mirrors (planes containing `axis`
    and bisecting between the two C2 ⟂ axes). One such plane is enough.
    """
    axis = _unit(axis)
    # Dn generators
    ops = build_Dn(n, axis=axis)
    # Add a dihedral mirror plane: rotate a vertical mirror by π/(2n)
    nrm = _vertical_mirror_normal(axis, angle=np.pi/(2*n))
    gens = ops + [mirror(nrm)]
    return group_closure(gens, max_elems=8*max(2, n))

def build_Sn(n, axis=(0,0,1)):
    """Pure improper rotation group generated by a single S_n (order n if n odd, 2n if n even)."""
    return group_closure([improper(axis, n)], max_elems=2*max(2, n))

def build_Cs(normal=(0,0,1)):
    """C_s: single mirror plane."""
    return group_closure([mirror(normal)], max_elems=2)

def build_Ci():
    """C_i: inversion only."""
    return group_closure([inversion()], max_elems=2)

# ---- Platonic families ----------------------------------------------------
def build_T():
    """
    Proper tetrahedral group T (order 12).
    Generators: 120° about (1,1,1) and 180° about (0,1,1) (twofold through opposite edges).
    """
    gens = [
        rot_axis_angle((1,1,1), 2*np.pi/3),
        rot_axis_angle((0,1,1), np.pi)
    ]
    return group_closure(gens, max_elems=12)

def build_Td():
    """
    Full tetrahedral group T_d (order 24).
    Add one improper S4 to T; closure generates σ_d and all S4 axes.
    """
    gens = [
        rot_axis_angle((1,1,1), 2*np.pi/3),
        rot_axis_angle((0,1,1), np.pi),
        improper((1,0,0), 4)
    ]
    return group_closure(gens, max_elems=24)

def build_Th():
    """
    T_h (order 24): T with inversion (no mirror planes).
    """
    gens = [
        rot_axis_angle((1,1,1), 2*np.pi/3),
        rot_axis_angle((0,1,1), np.pi),
        inversion()
    ]
    return group_closure(gens, max_elems=24)

def build_O():
    """
    Proper octahedral/cubic O (order 24).
    Generators: 90° about x; 120° about (1,1,1).
    """
    gens = [
        rot_axis_angle((1,0,0), np.pi/2),
        rot_axis_angle((1,1,1), 2*np.pi/3)
    ]
    return group_closure(gens, max_elems=24)

def build_Oh():
    """
    Full octahedral 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_I():
    """
    Proper icosahedral I (order 60).
    Practical generators: 5-fold about z and 3-fold about (1,1,1).
    """
    gens = [
        rot_axis_angle((0,0,1), 2*np.pi/5),
        rot_axis_angle((1,1,1), 2*np.pi/3),
    ]
    return group_closure(gens, max_elems=60)

def build_Ih():
    """
    Full icosahedral I_h (order 120).
    """
    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)



