import math
from icosphere import icosphere

def generate_mfem_sphere_mesh(subdivisions: int, outputFile: str):
    """Generate a P2 triangular sphere mesh and save in MFEM mesh v1.x format."""
    # 1. Generate a geodesic icosphere
    vertices, faces = icosphere(subdivisions)  # returns (Nverts×3) array and list of 3-tuples :contentReference[oaicite:14]{index=14}

    # 2. Build unique edges and mid-edge nodes
    edgeToMidIndex = {}
    edgesSet = set()
    for tri in faces:
        for i in range(3):
            v0, v1 = tri[i], tri[(i+1) % 3]
            edge = tuple(sorted((v0, v1)))
            edgesSet.add(edge)

    midPoints = []
    for edge in edgesSet:
        v0, v1 = edge
        x0, y0, z0 = vertices[v0]
        x1, y1, z1 = vertices[v1]
        mx, my, mz = (x0 + x1)/2, (y0 + y1)/2, (z0 + z1)/2
        r = math.sqrt(mx*mx + my*my + mz*mz)
        mx, my, mz = mx/r, my/r, mz/r  # project onto unit sphere :contentReference[oaicite:15]{index=15}
        edgeToMidIndex[edge] = len(vertices) + len(midPoints)
        midPoints.append((mx, my, mz))

    # 3. Assemble all node coordinates (vertices + mid-edge nodes)
    nodeCoords = [tuple(v) for v in vertices] + midPoints

    # 4. Write MFEM mesh file
    with open(outputFile, 'w') as f:
        # Header
        f.write("MFEM mesh v1.0\n\n")

        # Topology
        f.write("dimension\n2\n\n")
        f.write(f"elements\n{len(faces)}\n")
        for tri in faces:
            v0, v1, v2 = tri
            m0 = edgeToMidIndex[tuple(sorted((v0, v1)))]
            m1 = edgeToMidIndex[tuple(sorted((v1, v2)))]
            m2 = edgeToMidIndex[tuple(sorted((v2, v0)))]
            # attribute=1, type=2 (triangle), followed by 6 vertex indices for P2 :contentReference[oaicite:16]{index=16}
            f.write(f"1 2 {v0} {v1} {v2} {m0} {m1} {m2}\n")
        f.write("\n")

        # No boundary segments for a closed surface
        f.write("boundary\n0\n\n")

        # Vertices count (actual DOFs)
        f.write(f"vertices\n{len(nodeCoords)}\n\n")

        # Geometry: Nodes (P2 in 3D)
        f.write("nodes\n")
        f.write("FiniteElementSpace\n")
        f.write("FiniteElementCollection: H1_3D_P2\n")
        f.write("VDim: 3\n")
        f.write("Ordering: 1\n")
        for x, y, z in nodeCoords:
            f.write(f"{x} {y} {z}\n")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Generate a spherical mesh in MFEM format')
    parser.add_argument('--subdivisions', type=int, default=2, help='Number of subdivisions for the icosphere')
    parser.add_argument('--output', type=str, default='sphere.mesh', help='Output file name')
    args = parser.parse_args()

    # Generate the mesh
    generate_mfem_sphere_mesh(args.subdivisions, args.output)
    print(f"MFEM mesh generated: {args.output}")