import pygmsh
import meshio
import numpy as np
from math import sqrt

import argparse

def mesh_size_callback(dim, tag, x, y, z, lc, a, b, k):
    r = np.sqrt(x**2 + y**2 + z**2)
    size = a + (b-a) * (1-np.exp(-k*r))
    return size

def generate_spherical_mesh(radius=1.0, coreRes = 0.1, outRes = 0.1, k = 2):
    with pygmsh.occ.Geometry() as geom:
        originPoint = geom.add_point([0.0, 0.0, 0.0], mesh_size=coreRes)
        geom.add_ball([0.0, 0.0, 0.0], 1.0)

        geom.set_mesh_size_callback(lambda dim, tag, x, y, z, lc: mesh_size_callback(dim, tag, x, y, z, lc, coreRes, outRes, k))
        mesh = geom.generate_mesh(verbose=True)

    return mesh


def write_mfem_mesh(meshData, filename):
    meshio.write(filename, meshData, file_format='gmsh22')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate a spherical mesh')
    parser.add_argument('--radius', type=float, default=1.0, help='Radius of the sphere')
    parser.add_argument('--coreMeshRes', type=float, default=0.01, help='Core Mesh Resolution')
    parser.add_argument('--outerMeshRes', type=float, default=0.1, help='Outer Mesh Resolution')
    parser.add_argument('--k', type=float, default=2, help='Mesh size scaling factor')
    parser.add_argument('--output', type=str, default='sphere.msh', help='Output file name')
    args = parser.parse_args()

    meshData = generate_spherical_mesh(args.radius, args.coreMeshRes, args.outerMeshRes, args.k)
    write_mfem_mesh(meshData, args.output)