from pathlib import Path
import rich
from rich.syntax import Syntax
from math import ceil
from multiprocessing import cpu_count
import gmsh
import numpy as np

from ...helper.experiment import ExperimentConfig

console = rich.get_console()
print = console.print


DIE_PIXEL_SIZE = 8.434786e-3  # mm/pixel # 8.434786 um/pixel
DIE_RADIUS = 2.550608716  # mm # 2550.608716 um, 5_000 / (8.434786 * 2) + 6 pixels
DIE_HEIGHT = 8.51913386  # mm # 8519.13386 um, 1010 pixels
DIE_CENTER = [2.65695759, 2.65695759, 0]  # [2656.95759, 2656.95759, 0] um, 315, 315, 0 pixels


def set_diagnostic_options(experiment: ExperimentConfig, save_forces: int, save_strain_energy: int, save_swarm: int, save_solution: int,
                           save_diagnostics: int, save: bool) -> None:
    """
    Set diagnostic options for the experiment.

    :param experiment: The experiment configuration object.
    :param save_forces: Whether to save forces.
    :param save_swarm: Whether to save swarm data.
    :param save_solution: Whether to save solution data.
    :param save_diagnostics: Whether to save diagnostics.
    :param save_all: Whether to save all data.
    """
    if not save:
        return
    if save_forces > 0:
        experiment.diagnostic_options["ts_monitor_surface_force_per_face"] = "ascii:forces.csv"
        experiment.diagnostic_options["ts_monitor_surface_force_per_face_interval"] = f"{save_forces}"
    if save_strain_energy > 0:
        experiment.diagnostic_options["ts_monitor_strain_energy"] = "ascii:strain_energy.csv"
        experiment.diagnostic_options["ts_monitor_strain_energy_interval"] = f"{save_strain_energy}"
    if save_swarm > 0:
        experiment.diagnostic_options["ts_monitor_swarm_solution"] = "ascii:swarm.xmf"
        experiment.diagnostic_options["ts_monitor_swarm_fields"] = "J,volume,rho,material"
        experiment.diagnostic_options["ts_monitor_swarm_solution_interval"] = f"{save_swarm}"
    if save_solution > 0:
        experiment.diagnostic_options["ts_monitor_solution"] = r"cgns:solution_%06d.cgns"
        experiment.diagnostic_options["ts_monitor_solution_interval"] = f"{save_solution}"
    if save_diagnostics > 0:
        experiment.diagnostic_options["ts_monitor_output_fields"] = r"cgns:output_fields_%06d.cgns"
        experiment.diagnostic_options["ts_monitor_output_fields_interval"] = f"{save_diagnostics}"


def get_mesh(characteristic_length: float, voxel_data: Path,
             scratch_dir: Path, load_fraction=0.4, clamp_top=True) -> str:
    """
    Get a mesh file for the given voxel data and characteristic length, generating if necessary.

    :param characteristic_length: The desired characteristic length for the mesh.
    :param voxel_data: Path to the voxel data file (e.g., CT scan).
    :param scratch_dir: Directory to store the generated mesh file.
    :param load_fraction: Fraction of the load to apply to the die (default is 0.4).
    :param clamp_top: Whether to clamp the top of the die (default is True).

    :return: A dictionary of mesh options for the experiment configuration.
    """
    mesh_file: Path = generate_mesh(characteristic_length, voxel_data, scratch_dir)

    options: str = '\n'.join([
        "",
        "# Mesh options generated by press_common.generate_mesh",
        "mpm_voxel:",
        f"  filename: {voxel_data.resolve()}",
        f"  pixel_size: {DIE_PIXEL_SIZE}",
        "",
        "dm_plex:",
        f"  filename: {mesh_file.resolve()}",
        "  cgns_parallel:",
        "  box_label:",
        "  dim: 3",
        "  simplex: 0",
        "",
        "remap:",
        f"  direction: z",
        f"  scale: {(1 - load_fraction)} # (1 - load_fraction) to match displacement",
        "",
    ])
    if clamp_top:
        options += '\n'.join([
            "bc:",
            "  clamp: 1,2",
            "  # Clamped displacement for top and bottom",
            "  clamp_2:",
            f"    translate: 0,0,{-load_fraction * DIE_HEIGHT} # -load_fraction * height",
            "  # Prevent x,y expansion beyond the die boundary",
            "  slip: 3,4,5,6",
            "  slip_3:",
            "    components: 0,1",
            "  slip_4:",
            "    components: 0,1",
            "  slip_5:",
            "    components: 0,1",
            "  slip_6:",
            "    components: 0,1",
            "",
        ])
    else:
        options += '\n'.join([
            "bc:",
            "  slip: 1,2,3,4,5,6",
            "  # Allow x,y displacement for top and bottom, prescribe z displacement",
            "  slip_1:",
            "    components: 2",
            "  slip_2:",
            "    components: 2",
            f"    translate: {-load_fraction * DIE_HEIGHT} # -load_fraction * height",
            "  slip_3:",
            "    components: 0,1",
            "  slip_4:",
            "    components: 0,1",
            "  slip_5:",
            "    components: 0,1",
            "  slip_6:",
            "    components: 0,1",
            "",
        ])

    print(f"[info]Generated mesh options:[/]")
    syntax = Syntax(options, "yaml")
    print(syntax)

    return options


def generate_mesh(characteristic_length: float, voxel_data: Path, scratch_dir: Path) -> Path:
    mesh_dir = scratch_dir / "meshes"
    if not mesh_dir.exists():
        mesh_dir.mkdir(parents=True, exist_ok=True)
    if not voxel_data.exists():
        raise FileNotFoundError(f"Voxel data {voxel_data} does not exist")
    mesh_file = mesh_dir / f"cylinder_{voxel_data.stem}_CL{characteristic_length}.cgns"
    if mesh_file.exists():
        console.print(f"[info]Using existing mesh [/]{mesh_file}")
        return mesh_file.resolve()

    element_order = 1
    center = DIE_CENTER
    radius = DIE_RADIUS
    height = DIE_HEIGHT
    square_radius = DIE_RADIUS * np.sqrt(0.125)

    layers = int(np.ceil(DIE_HEIGHT / characteristic_length) + 1)
    print(f"[info]Center: [/]{center}")
    print(f"[info]Radius: [/]{radius}")
    print(f"[info]Height: [/]{height}")
    print(f"[info]Number of layers: [/]{layers}")
    print(f"[info]Element order: [/]{element_order}")
    print(f"[info]Characteristic length: [/]{characteristic_length}")

    EDGE = 1
    SURFACE = 2
    VOLUME = 3

    inner_rad = radius - 0.5 * square_radius * (1 + np.sqrt(2))
    outer_rad = radius - square_radius
    avg_rad = 0.5 * (inner_rad + outer_rad)
    NPTS_SQUARE = int(np.ceil(0.5 / characteristic_length * (square_radius + np.pi / 4 * radius))) + 1
    NPTS_RADIAL = int(np.ceil(avg_rad / characteristic_length)) + 1

    gmsh.initialize()

    # set mesh options
    gmsh.option.setNumber("General.NumThreads", cpu_count())
    gmsh.option.setNumber("Mesh.Algorithm3D", 10)
    gmsh.option.setNumber("Mesh.ElementOrder", element_order)
    gmsh.option.setNumber("Mesh.HighOrderOptimize", 1 if element_order > 1 else 0)
    gmsh.option.setNumber("Mesh.Binary", 1)

    p1 = gmsh.model.geo.addPoint(*center)
    p2 = gmsh.model.geo.addPoint(*(center + np.array([square_radius, 0, 0])))
    offset = (0.5 + 0.5 / np.sqrt(2))
    p3 = gmsh.model.geo.addPoint(*(center + np.array([square_radius * offset, square_radius * offset, 0])))
    p4 = gmsh.model.geo.addPoint(*(center + np.array([0, square_radius, 0])))

    p5 = gmsh.model.geo.addPoint(*(center + np.array([radius, 0, 0])))
    p6 = gmsh.model.geo.addPoint(*(center + np.array([radius / np.sqrt(2), radius / np.sqrt(2), 0])))
    p7 = gmsh.model.geo.addPoint(*(center + np.array([0, radius, 0])))

    l1 = gmsh.model.geo.addLine(p1, p2)
    l2 = gmsh.model.geo.addLine(p2, p3)
    l3 = gmsh.model.geo.addLine(p3, p4)
    l4 = gmsh.model.geo.addLine(p4, p1)

    l5 = gmsh.model.geo.addLine(p2, p5)
    l6 = gmsh.model.geo.addCircleArc(p5, p1, p6)
    l7 = gmsh.model.geo.addLine(p6, p3)
    l8 = gmsh.model.geo.addCircleArc(p6, p1, p7)
    l9 = gmsh.model.geo.addLine(p7, p4)

    cl1 = gmsh.model.geo.addCurveLoop((l1, l2, l3, l4))
    cl2 = gmsh.model.geo.addCurveLoop((l5, l6, l7, -l2))
    cl3 = gmsh.model.geo.addCurveLoop((l9, -l3, -l7, l8))

    s1 = gmsh.model.geo.addPlaneSurface([cl1])
    s2 = gmsh.model.geo.addPlaneSurface([cl2])
    s3 = gmsh.model.geo.addPlaneSurface([cl3])

    for c in [l5, l7, l9]:
        gmsh.model.geo.mesh.setTransfiniteCurve(c, NPTS_RADIAL)
    for c in [l1, l2, l3, l4, l6, l8]:
        gmsh.model.geo.mesh.setTransfiniteCurve(c, NPTS_SQUARE)
    for s in [s1, s2, s3]:
        gmsh.model.geo.mesh.setTransfiniteSurface(s)
        gmsh.model.geo.mesh.setRecombine(SURFACE, s)

    gmsh.model.geo.synchronize()

    def add_quadrant(orig, angle):
        quadrant = gmsh.model.geo.copy(orig)
        lines = [id for dim, id in quadrant if dim == EDGE]
        for l in lines:
            if l - lines[0] + 1 in [l5, l7, l9]:
                gmsh.model.geo.mesh.setTransfiniteCurve(l, NPTS_RADIAL)
            else:
                gmsh.model.geo.mesh.setTransfiniteCurve(l, NPTS_SQUARE)
        for s in filter(lambda x: x[0] == SURFACE, quadrant):
            gmsh.model.geo.mesh.setTransfiniteSurface(s[1])
            gmsh.model.geo.mesh.setRecombine(s[0], s[1])
        gmsh.model.geo.rotate(quadrant, center[0], center[1], center[2], 0, 0, 1, angle)
        return list(set(gmsh.model.getEntities()).intersection(quadrant))

    quadrant1 = gmsh.model.getEntities()
    add_quadrant(quadrant1, np.pi / 2)
    add_quadrant(quadrant1, np.pi)
    add_quadrant(quadrant1, 3 * np.pi / 2)

    gmsh.model.geo.removeAllDuplicates()
    gmsh.model.geo.synchronize()

    disk = gmsh.model.getEntities(SURFACE)

    # Extrude the disk along (0, 0, height) with the specified number of layers.
    # The 'recombine=True' option tells gmsh to recombine triangular faces into quads.
    gmsh.model.geo.extrude(disk, 0, 0, height, [layers], recombine=True)

    # Synchronize to update the model with the new entities.
    gmsh.model.geo.synchronize()

    print('About to generate mesh')

    # generate mesh
    gmsh.model.mesh.generate(3)

    node_tags, coords, parametric_coords = gmsh.model.mesh.getNodes()
    types, tags, nodes = gmsh.model.mesh.getElements(VOLUME)
    gmsh.finalize()

    gmsh.initialize()
    v1 = gmsh.model.addDiscreteEntity(VOLUME)
    gmsh.model.mesh.addNodes(VOLUME, v1, node_tags, coords, parametric_coords)
    gmsh.model.mesh.addElements(VOLUME, v1, types, tags, nodes)
    # save mesh
    gmsh.write(str(mesh_file))
    gmsh.finalize()

    if not mesh_file.exists():
        raise RuntimeError(f"Failed to generate mesh file: {mesh_file}")

    print(f"[info]Mesh saved to [/]{mesh_file}")
    return mesh_file
