# EMerge is an open source Python based FEM EM simulation module.
# Copyright (C) 2025  Robert Fennis.

# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program; if not, see
# <https://www.gnu.org/licenses/>.

from typing import TypeVar
from ..geometry import GeoSurface, GeoVolume
from ..cs import CoordinateSystem, GCS
import gmsh
import numpy as np

T = TypeVar('T', GeoSurface, GeoVolume)

def _gen_mapping(obj_in, obj_out) -> dict:
    tag_mapping: dict[int, dict] = {0: dict(),
                                        1: dict(),
                                        2: dict(),
                                        3: dict()}
    for domain, mapping in zip(obj_in, obj_out):
        tag_mapping[domain[0]][domain[1]] = [o[1] for o in mapping]
    return tag_mapping

def add(main: T, tool: T, 
             remove_object: bool = True,
             remove_tool: bool = True) -> T:
    ''' Adds two GMSH objects together, returning a new object that is the union of the two.
    
    Parameters
    ----------
    main : GeoSurface | GeoVolume
    tool : GeoSurface | GeoVolume
    remove_object : bool, optional
        If True, the main object will be removed from the model after the operation. Default is True.
    remove_tool : bool, optional
        If True, the tool object will be removed from the model after the operation. Default is True.
    
    Returns
    -------
    GeoSurface | GeoVolume
        A new object that is the union of the main and tool objects.
    '''
    out_dim_tags, out_dim_tags_map = gmsh.model.occ.fuse(main.dimtags, tool.dimtags, removeObject=remove_object, removeTool=remove_tool)
    if out_dim_tags[0][0] == 3:
        output = GeoVolume([dt[1] for dt in out_dim_tags])._take_tools(tool,main)
    elif out_dim_tags[0][0] == 2:
        output = GeoSurface([dt[1] for dt in out_dim_tags])._take_tools(tool,main)
    if remove_object:
        main._exists = False
    if remove_tool:
        tool._exists = False
    return output

def remove(main: T, tool: T, 
             remove_object: bool = True,
             remove_tool: bool = True) -> T:
    ''' Subtractes a tool object GMSH from the main object, returning a new object that is the difference of the two.
    
    Parameters
    ----------
    main : GeoSurface | GeoVolume
    tool : GeoSurface | GeoVolume
    remove_object : bool, optional
        If True, the main object will be removed from the model after the operation. Default is True.
    remove_tool : bool, optional
        If True, the tool object will be removed from the model after the operation. Default is True.
    
    Returns
    -------
    GeoSurface | GeoVolume
        A new object that is the difference of the main and tool objects.
    '''
    out_dim_tags, out_dim_tags_map = gmsh.model.occ.cut(main.dimtags, tool.dimtags, removeObject=remove_object, removeTool=remove_tool)
    if out_dim_tags[0][0] == 3:
        output = GeoVolume([dt[1] for dt in out_dim_tags])._take_tools(tool,main)
    elif out_dim_tags[0][0] == 2:
        output = GeoSurface([dt[1] for dt in out_dim_tags])._take_tools(tool,main)
    if remove_object:
        main._exists = False
    if remove_tool:
        tool._exists = False
    return output

subtract = remove

def intersect(main: T, tool: T, 
             remove_object: bool = True,
             remove_tool: bool = True) -> T:
    ''' Intersection of a tool object GMSH with the main object, returning a new object that is the intersection of the two.
    
    Parameters
    ----------
    main : GeoSurface | GeoVolume
    tool : GeoSurface | GeoVolume
    remove_object : bool, optional
        If True, the main object will be removed from the model after the operation. Default is True.
    remove_tool : bool, optional
        If True, the tool object will be removed from the model after the operation. Default is True.
    
    Returns
    -------
    GeoSurface | GeoVolume
        A new object that is the difference of the main and tool objects.
    '''
    out_dim_tags, out_dim_tags_map = gmsh.model.occ.intersect(main.dimtags, tool.dimtags, removeObject=remove_object, removeTool=remove_tool)
    if out_dim_tags[0][0] == 3:
        output = GeoVolume([dt[1] for dt in out_dim_tags])._take_tools(tool,main)
    elif out_dim_tags[0][0] == 2:
        output = GeoSurface([dt[1] for dt in out_dim_tags])._take_tools(tool,main)
    if remove_object:
        main._exists = False
    if remove_tool:
        tool._exists = False
    return output

def embed(main: GeoVolume, other: GeoSurface) -> None:
    ''' Embeds a surface into a volume in the GMSH model.
    Parameters
    ----------
    main : GeoVolume
        The volume into which the surface will be embedded.
    other : GeoSurface
        The surface to be embedded into the volume.
    
    Returns
    -------
    None
    '''
    gmsh.model.geo.synchronize()
    gmsh.model.mesh.embed(other.dim, [other.tag,], main.dim, main.tags)

def rotate(main: GeoVolume, 
           c0: tuple[float, float, float],
           ax: tuple[float, float, float],
           angle: float,
           degree=True) -> GeoVolume:
    """Rotates a GeoVolume object around an axist defined at a coordinate.

    Args:
        main (GeoVolume): The object to rotate
        c0 (tuple[float, float, float]): The point of origin for the rotation axis
        ax (tuple[float, float, float]): A vector defining the rotation axis
        angle (float): The angle in degrees (if degree is True)
        degree (bool, optional): Whether to interpret the angle in degrees. Defaults to True.

    Returns:
        GeoVolume: The rotated GeoVolume object.
    """
    if degree:
        angle = angle * np.pi/180
    gmsh.model.occ.rotate(main.dimtags, *c0, *ax, -angle)

    # Rotate the facepointers
    for fp in main._all_pointers:
        fp.rotate(c0, ax, angle)
    return main

def translate(main: GeoVolume,
              dx: float = 0,
              dy: float = 0,
              dz: float = 0) -> GeoVolume:
    """Translates the GeoVolume object along a given displacement

    Args:
        main (GeoVolume): The object to translate
        dx (float, optional): The X-displacement in meters. Defaults to 0.
        dy (float, optional): The Y-displacement in meters. Defaults to 0.
        dz (float, optional): The Z-displacement in meters. Defaults to 0.

    Returns:
        GeoVolume: The translated object
    """
    gmsh.model.occ.translate(main.dimtags, dx, dy, dz)
    
     # Rotate the facepointers
    for fp in main._all_pointers:
        fp.translate(dx, dy, dz)

    return main

def mirror(main: GeoVolume,
           origin: tuple[float, float, float] = (0.0, 0.0, 0.0),
           direction: tuple[float, float, float] = (1.0, 0.0, 0.0),
           make_copy: bool = True) -> GeoVolume:
    """Mirrors a GeoVolume object along a miror plane defined by a direction originating at a point

    Args:
        main (GeoVolume): The object to mirror
        origin (tuple[float, float, float], optional): The point of origin in meters. Defaults to (0.0, 0.0, 0.0).
        direction (tuple[float, float, float], optional): The normal axis defining the plane of reflection. Defaults to (1.0, 0.0, 0.0).

    Returns:
        GeoVolume: The mirrored GeoVolume object
    """
    a, b, c = direction
    x0, y0, z0 = origin
    d = -(a*x0 + b*y0 + c*z0)
    if (a==0) and (b==0) and (c==0):
        return main
    mirror_obj = main
    if make_copy:
        new_obj = main.make_copy()
        gmsh.model.occ.mirror(new_obj.dimtags, a,b,c,d)
        mirror_obj = new_obj
    else:
        gmsh.model.occ.mirror(main.dimtags, a,b,c,d)

    for fp in mirror_obj._all_pointers:
        fp.mirror(origin, direction)
    return mirror_obj

def change_coordinate_system(main: GeoVolume,
                             new_cs: CoordinateSystem = GCS,
                             old_cs: CoordinateSystem = GCS):
    """Moves the GeoVolume object from a current coordinate system to a new one.

    The old and new coordinate system by default are the global coordinate system.
    Thus only one needs to be provided to transform to and from these coordinate systems.

    Args:
        main (GeoVolume): The object to transform
        new_cs (CoordinateSystem): The new coordinate system. Defaults to GCS
        old_cs (CoordinateSystem, optional): The old coordinate system. Defaults to GCS.

    Returns:
        _type_: _description_
    """
    if new_cs._is_global and old_cs._is_global:
        return main
    
    M1 = old_cs.affine_to_global()
    M2 = new_cs.affine_from_global()
    # Transform to the global coordinate system.
    if not old_cs._is_global:
        gmsh.model.occ.affine_transform(main.dimtags, M1.flatten()[:12])
    # Transform to a new coordinate system.
    if not new_cs._is_global:
        gmsh.model.occ.affineTransform(main.dimtags, M2.flatten()[:12])

    for fp in main._all_pointers:
        fp.affine_transform(M1)
        fp.affine_transform(M2)
    return main