#!python
import numpy as np

def main( geom, mesh ):
  cubit.cmd( "reset" )
  cubit.cmd( "set duplicate block elements off" )
  cubit.cmd( "set node constraint on" )

  # Extract variables
  a = geom["a"]
  b = geom["b"]
  c = geom["c"]
  z_cav = geom["z_cav"]
  rot_x = geom["rot_x"]
  rot_y = geom["rot_y"]
  rot_z = geom["rot_z"]
  pad = geom["pad"]
  domain_x = geom["domain_x"]
  domain_y = geom["domain_y"]
  domain_z = geom["domain_z"]
  cavity_region_element_size = mesh["cavity_region_element_size"]
  default_element_size = mesh["default_element_size"]

  # Outer computational domain
  cubit.cmd( f"create brick x {domain_x} y {domain_y} z {domain_z}" )
  cubit.cmd( f"volume 1 name 'outer_domain'" )
  cubit.cmd( f"move volume 1 z {-1 * domain_z / 2}" )

  # Ellipsoidal cavity tool
  cubit.cmd( f"create sphere radius 1" )
  cubit.cmd( f"split periodic volume all" )
  cubit.cmd( f"volume 2 name 'cavity_tool'" )
  cubit.cmd( f"volume 2 scale x {a} y {b} z {c}" )

  # Rotation
  cubit.cmd( f"volume 2 rotate {rot_z} about z" )
  cubit.cmd( f"volume 2 rotate {rot_y} about y" )
  cubit.cmd( f"volume 2 rotate {rot_x} about x" )

  cubit.cmd( f"move volume 2 x 0 y 0 z {z_cav}" )

  # Boolean subtraction to create the void
  cubit.cmd( f"subtract volume 2 from volume 1" )

  # Split domain for meshing control
  cubit.cmd( f"webcut volume all with plane xplane offset 0" )
  cubit.cmd( f"webcut volume all with plane yplane offset 0" )
  cubit.cmd( f"webcut volume all with plane zplane offset {z_cav}" )
  cubit.cmd( f"webcut volume all with plane zplane offset {2*z_cav}" )

  major_radius, minor_radius, rotate_about_z = projected_ellipse_for_cubit( a, b, c, rot_x, rot_y, rot_z, pad )

  cubit.cmd( f"create surface ellipse major radius {major_radius} minor radius {minor_radius} zplane" )
  surface_id = cubit.get_last_id( "surface" )
  cubit.cmd( f"surface {surface_id} rotate {rotate_about_z} about z" )

  cubit.cmd( f"sweep curve {cubit.get_last_id('curve')}  vector 0 0 -1  distance 100000 keep" )
  cubit.cmd( f"webcut volume all with sheet extended from surface {cubit.get_last_id('surface')}" )
  cubit.cmd( f"delete surface all" )

  ellipse_surfs = find_all_spline_surfaces()
  print( ellipse_surfs )
  cubit.cmd( f"composite create surface {cubit.get_id_string(ellipse_surfs)} keep angle 15" )

  cubit.cmd( "merge all" )

  # External boundaries
  cubit.cmd( f"sideset 1 add surface with x_coord <= -{domain_x/2};   sideset 1 name 'xmin'" )
  cubit.cmd( f"sideset 2 add surface with x_coord >= {domain_x/2};    sideset 2 name 'xmax'" )
  cubit.cmd( f"sideset 3 add surface with y_coord <= -{domain_y/2};   sideset 3 name 'ymin'" )
  cubit.cmd( f"sideset 4 add surface with y_coord >= {domain_y/2};    sideset 4 name 'ymax'" )
  cubit.cmd( f"sideset 5 add surface with z_coord <= -{domain_z - 1}; sideset 5 name 'zmin'" )
  cubit.cmd( f"sideset 6 add surface with z_coord >= -1;              sideset 6 name 'zmax'" )

  # Cavity surface extraction (curvature filter avoids picking webcut planes)
  cubit.cmd( f'group "cavity_surfs" add surface {cubit.get_id_string(ellipse_surfs)}' )
  cubit.cmd( f"sideset 99 add surface in group cavity_surfs" )
  cubit.cmd( f"sideset 99 name 'magma_boundary'" )

  # Node Sets
  cubit.cmd( "nodeset 1 surface in sideset 1" )
  cubit.cmd( "nodeset 2 surface in sideset 2" )
  cubit.cmd( "nodeset 3 surface in sideset 3" )
  cubit.cmd( "nodeset 4 surface in sideset 4" )
  cubit.cmd( "nodeset 5 surface in sideset 5" )
  cubit.cmd( "nodeset 6 surface in sideset 6" )
  cubit.cmd( "nodeset 99 surface in sideset 99" )
  
  cubit.cmd( "reset volume all" )
  
  cubit.cmd( f"volume all size {default_element_size}" )

  poly_vols = find_polyhedron_scheme_volumes()
  cubit.cmd( f"group 'poly_vols' add Volume {cubit.get_id_string( poly_vols )}" )

  cubit.cmd( f"volume in poly_vols scheme polyhedron" )
  cubit.cmd( f"volume in poly_vols size {cavity_region_element_size}" )
  cubit.cmd( f"mesh volume in poly_vols" )
  cubit.cmd( f"volume all except in poly_vols size {default_element_size}" )

  cubit.cmd( "group 'sweep_vols' add volume all except in poly_vols" )

  group_id = get_group_id_from_name( "sweep_vols" )
  for vid in cubit.get_group_volumes( group_id ):
    sweep_surf_ids = find_volume_sweep_surfs( vid )
    source_id = sweep_surf_ids["source"]
    target_id = sweep_surf_ids["target"]
    cubit.cmd( f"surface {source_id} scheme pave" )
    cubit.cmd( f"volume {vid} redistribute nodes off" )
    cubit.cmd( f"volume {vid} autosmooth target on  fixed imprints off  smart smooth off" )
    cubit.cmd( f"volume {vid} scheme Sweep  source surface {source_id} target surface {target_id} sweep transform least squares" )
    cubit.cmd( f"mesh volume {vid}" )
  
  # Smoothing
  cubit.cmd( "volume all smooth scheme condition number beta 2.0 cpu 30" )
  cubit.cmd( "smooth volume all" )

  # Export
  cubit.cmd( "export abaqus 'ellipsoid_cavity.inp' overwrite" )
  cubit.cmd( "save cub5 'ellipsoid_cavity.cub5' overwrite" )

    
def normalize_180(angle_deg):
  """
  Normalize an angle to [-90, 90) degrees.

  For an ellipse, rotations that differ by 180 degrees are equivalent.
  """
  return float((angle_deg + 90.0) % 180.0 - 90.0)


def cubit_rotation_matrix(rot_x, rot_y, rot_z):
  """
  Return the rotation matrix for the Cubit command sequence:

    rotate rot_z about z
    rotate rot_y about y
    rotate rot_x about x

  Angles are in degrees.

  With column vectors, this corresponds to:

    R = Rx @ Ry @ Rz
  """

  rx, ry, rz = np.radians([rot_x, rot_y, rot_z])

  cx, sx = np.cos(rx), np.sin(rx)
  cy, sy = np.cos(ry), np.sin(ry)
  cz, sz = np.cos(rz), np.sin(rz)

  Rx = np.array([
    [1.0, 0.0, 0.0],
    [0.0, cx, -sx],
    [0.0, sx, cx],
  ])

  Ry = np.array([
    [cy, 0.0, sy],
    [0.0, 1.0, 0.0],
    [-sy, 0.0, cy],
  ])

  Rz = np.array([
    [cz, -sz, 0.0],
    [sz, cz, 0.0],
    [0.0, 0.0, 1.0],
  ])

  return Rx @ Ry @ Rz


def projected_ellipse_for_cubit(a, b, c, rot_x, rot_y, rot_z, pad=0.0):
  """
  Compute the ellipse obtained by orthogonally projecting a rotated ellipsoid
  onto the xy-plane.

  The ellipsoid starts as a unit sphere scaled by:

    x -> a
    y -> b
    z -> c

  and is then rotated using the Cubit command sequence:

    rotate rot_z about z
    rotate rot_y about y
    rotate rot_x about x

  Returns
  -------
  major_radius : float
    Padded semimajor radius of the projected ellipse.

  minor_radius : float
    Padded semiminor radius of the projected ellipse.

  rotate_about_z : float
    Rotation angle, in degrees, for the Cubit zplane ellipse.
    This aligns the local x-axis of the ellipse with the semimajor axis.
  """

  R = cubit_rotation_matrix(rot_x, rot_y, rot_z)

  # Squared semiaxis matrix for the original, unrotated ellipsoid.
  D2 = np.diag([a*a, b*b, c*c])

  # Projection from xyz to xy.
  P = np.array([
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 0.0],
  ])

  # The projected ellipse is described by:
  #
  #   q.T @ inv(K) @ q <= 1
  #
  # where q = [x, y].
  K = P @ R @ D2 @ R.T @ P.T

  # Eigenvalues of K are the squared semiaxis lengths.
  # Eigenvectors give the corresponding directions in the xy-plane.
  eigvals, eigvecs = np.linalg.eig(K)

  # np.linalg.eigh returns eigenvalues in ascending order.
  # Sort them from largest to smallest.
  order = np.argsort(eigvals)[::-1]
  eigvals = eigvals[order]
  eigvecs = eigvecs[:, order]

  major_radius = np.sqrt(eigvals[0]) + pad
  minor_radius = np.sqrt(eigvals[1]) + pad

  # Direction of the semimajor axis in the xy-plane.
  major_direction = eigvecs[:, 0]

  rotate_about_z = np.degrees(
    np.arctan2(major_direction[1], major_direction[0])
  )

  rotate_about_z = normalize_180(rotate_about_z)

  return float(major_radius), float(minor_radius), float(rotate_about_z)

def find_all_spline_surfaces():
  spline_surfs = []
  S = cubit.get_entities( "surface" )
  for sid in S:
    if cubit.get_surface_type( sid ) == "spline surface":
      spline_surfs.append( sid )
  return spline_surfs

def find_polyhedron_scheme_volumes():
  poly_vols = []
  V = cubit.get_entities( "volume" )
  for vid in V:
    S = cubit.parse_cubit_list( "surface", f"in volume {vid}" )
    for sid in S:
      if cubit.get_surface_type( sid ) == "spline surface":
        poly_vols.append( vid )
        continue
      elif cubit.is_virtual( "surface", sid ):
        poly_vols.append( vid )
        continue
  return poly_vols

def find_volume_sweep_surfs( vid ):
  S = cubit.parse_cubit_list( "surface", f"in volume {vid}" )
  zmax = -1 * float( "inf" )
  zmin = float( "inf" )
  for sid in S:
    x,y,z = cubit.surface( sid ).center_point()
    if z > zmax:
      zmax = z
      zmax_surf = sid
    if z < zmin:
      zmin = z
      zmin_surf = sid
  return {"source": zmax_surf, "target": zmin_surf}

def get_group_id_from_name( name ):
  groups = cubit.group_names_ids()
  for group in groups:
    if group[0] == name:
      return group[1]
  return None

if __name__ == "__coreformcubit__":
  geom = {}
  # Ellipsoid semi-axes (meters)
  geom["a"] = 2000
  geom["b"] = 2710
  geom["c"] = 3032
  geom["z_cav"] = -6860
  geom["pad"] = 2000

  # Rotation angle (degrees)
  geom["rot_x"] = 30
  geom["rot_y"] = 20
  geom["rot_z"] = 10

  geom["domain_x"] = 100000
  geom["domain_y"] = 100000
  geom["domain_z"] = 100000

  mesh = {}
  mesh["cavity_region_element_size"] = 400
  mesh["default_element_size"] = 2000
  main( geom, mesh )