# This script generates a 2D Gmsh geometry for the ANASEN detector's wire planes, including an optional hot needle and IC wires. The geometry is designed for adaptive meshing, with finer mesh near the wires and coarser mesh in empty space. The script takes a z-locus as a command-line argument to determine the interpolation between the two wire planes.
import numpy as np
import gmsh, sys

# --- Configuration Flags ---
include_ic_wires = True
include_needle = True

gmsh.initialize()
gmsh.model.add("adaptive_mesh")
gmsh.option.setNumber("General.NumThreads", 10)
# gmsh.option.setNumber("Mesh.Adapt.MaxNumberOfElements", 200000)
# gmsh.option.setNumber("Mesh.Adapt.MaxNumberOfNodes", 200000)
# gmsh.option.setNumber("Mesh.Adapt.MaxIter",5)
# gmsh.option.setNumber("Mesh.MeshSizeMin", 5e-3)
# gmsh.option.setNumber("Mesh.MeshSizeMax", 10.0)
gmsh.option.setNumber("Geometry.Tolerance", 4e-2)
# gmsh.option.setNumber("Mesh.MeshSizeExtendFromBoundary", 0)
# Only needed alongside mesh.recombine(), which is off. See the note there.
# gmsh.option.setNumber("Mesh.RecombinationAlgorithm", 2)
# gmsh.option.setNumber("Mesh.SubdivisionAlgorithm", 1)
gmsh.option.setNumber("Mesh.SecondOrderIncomplete", 1)

lc = 0.04
# z_loc = -174.3

if len(sys.argv) < 2:
    print("Usage: python3 wires_gmsh2d_bc.py <z_locus in mm> [selected_cathode_index]")
    quit()

z_loc = float(sys.argv[1])

# Cathode wire (0..23) given its own physical group, tag 40, so that
# wires2d_weight.sif can hold it at 1 V while everything else stays grounded.
selected_cathode = int(sys.argv[2]) if len(sys.argv) > 2 else 1

wireShift = 4.0
k = 2 * np.pi / 24.0
kg = k/2.0

# Plane 1 Offsets (-zmax/2)
# Anodes: -6*k (base) - 4*k (correction) = -10*k
offset_a1 = -6 * k - 4 * k
# Cathodes: -6*k (base) + 0.5*k (half-placement)
offset_c1 = -6 * k + (k / 2.0)
# Guard wires: aligned with cathodes
offset_g1 = offset_c1

# Plane 2 Offsets (+zmax/2) with Twist
offset_a2 = offset_a1 + (wireShift * k)
offset_c2 = offset_c1 - (wireShift * k)
offset_g2 = offset_c2

# 1 needle, 24 ic1, 24 ic2, 48 guard wires, 24 anodes, 24 cathodes

# needle at plane 1 at -zmax/2 no rotation
xarr_needle = np.array([0])
yarr_needle = np.array([0])

# ic1 wires, plane 1 at -zmax/2 no rotation
xarr_i11 = np.array([23 * np.cos(k * i) for i in range(24)])
yarr_i11 = np.array([23 * np.sin(k * i) for i in range(24)])

# ic1 wires, plane 1 at -zmax/2 no rotation
xarr_i21 = np.array([23 * np.cos(k * i + k / 2.0) for i in range(24)])
yarr_i21 = np.array([23 * np.sin(k * i + k / 2.0) for i in range(24)])

# --- Coordinate Arrays (Plane 1) ---
# Anodes: -k*i (Left-handed twist direction)
xarra_1 = np.array([37 * np.cos(-k * i + offset_a1) for i in range(24)])
yarra_1 = np.array([37 * np.sin(-k * i + offset_a1) for i in range(24)])

# Cathodes: +k*i (Right-handed twist direction)
xarrc_1 = np.array([42 * np.cos(k * i + offset_c1) for i in range(24)])
yarrc_1 = np.array([42 * np.sin(k * i + offset_c1) for i in range(24)])

# Guard Wires (48 wires, use kg spacing)
xarrg_1 = np.array([32 * np.cos(kg * i + offset_g1) for i in range(48)])
yarrg_1 = np.array([32 * np.sin(kg * i + offset_g1) for i in range(48)])

# needle at plane 2 at zmax/2
xarr_needle_2 = np.array([0])
yarr_needle_2 = np.array([0])

# #ic1 wires, plane 2 at zmax/2
xarr_i12 = np.array([23 * np.cos(k * i) for i in range(24)])
yarr_i12 = np.array([23 * np.sin(k * i) for i in range(24)])

# #ic2 wires, plane 2 at zmax/2
xarr_i22 = np.array([23 * np.cos(k * i + k / 2.0) for i in range(24)])
yarr_i22 = np.array([23 * np.sin(k * i + k / 2.0) for i in range(24)])

# --- Coordinate Arrays (Plane 2) ---
xarra_2 = np.array([37 * np.cos(-k * i + offset_a2) for i in range(24)])
yarra_2 = np.array([37 * np.sin(-k * i + offset_a2) for i in range(24)])

xarrc_2 = np.array([42 * np.cos(k * i + offset_c2) for i in range(24)])
yarrc_2 = np.array([42 * np.sin(k * i + offset_c2) for i in range(24)])

xarrg_2 = np.array([32 * np.cos(kg * i + offset_g2) for i in range(48)])
yarrg_2 = np.array([32 * np.sin(kg * i + offset_g2) for i in range(48)])

direction_needle_x = xarr_needle_2 - xarr_needle
direction_needle_y = yarr_needle_2 - yarr_needle

direction_ic1_x = xarr_i12 - xarr_i11
direction_ic1_y = yarr_i12 - yarr_i11

direction_ic2_x = xarr_i22 - xarr_i21
direction_ic2_y = yarr_i22 - yarr_i21

direction_guard_x = xarrg_2 - xarrg_1
direction_guard_y = yarrg_2 - yarrg_1

direction_anodes_x = xarra_2 - xarra_1
direction_anodes_y = yarra_2 - yarra_1

direction_cathodes_x = xarrc_2 - xarrc_1
direction_cathodes_y = yarrc_2 - yarrc_1

t = (z_loc + 174.3) / (2 * 174.3)  # z=-174.3 is 0, z=+174.3 is 1

xloc_needle = xarr_needle + t * direction_needle_x
yloc_needle = yarr_needle + t * direction_needle_y
xloc_i1 = xarr_i11 + t * direction_ic1_x
yloc_i1 = yarr_i11 + t * direction_ic1_y
xloc_i2 = xarr_i21 + t * direction_ic2_x
yloc_i2 = yarr_i21 + t * direction_ic2_y
xloc_g = xarrg_1 + t * direction_guard_x
yloc_g = yarrg_1 + t * direction_guard_y
xloc_a = xarra_1 + t * direction_anodes_x
yloc_a = yarra_1 + t * direction_anodes_y
xloc_c = xarrc_1 + t * direction_cathodes_x
yloc_c = yarrc_1 + t * direction_cathodes_y

# wire_radius_a = 0.018 #mm
# wire_radius_c = 0.0762 #mm
# wire_radius_g = 0.0762 #mm
wire_radius = 0.254  # mm
needle = []
ic1_wires = []
ic2_wires = []
guard_wires = []
anode_wires = []
cathode_wires = []
iw1_tags = [(3, i) for i in range(24)]
iw2_tags = [(3, i + 24) for i in range(24)]
gw_tags = [(3, i + 48) for i in range(48)]
aw_tags = [(3, i) for i in range(24)]
cw_tags = [(3, i + 24) for i in range(24)]

# for i,[xa,ya,xc,yc] in enumerate(zip(xarra_1,yarra_1,xarrc_1,yarrc_1)):
# create Hot Needle (1 total)
for i, (xn, yn) in enumerate(zip(xloc_needle, yloc_needle)):
    if include_needle:
        ndisk = gmsh.model.occ.addDisk(xn, yn, 0, wire_radius, wire_radius)
        needle.append(ndisk)

# create Guard Wires (48 total)
for xg, yg in zip(xloc_g, yloc_g):
    gdisk = gmsh.model.occ.addDisk(xg, yg, 0, wire_radius, wire_radius)
    guard_wires.append(gdisk)

# create Cathode Wires (24 total)
for xc, yc in zip(xloc_c, yloc_c):
    cdisk = gmsh.model.occ.addDisk(xc, yc, 0, wire_radius, wire_radius)
    cathode_wires.append(cdisk)

# create IC Anode and Cathode Wires (24 total each)
for i, (xa, ya) in enumerate(zip(xloc_a, yloc_a)):
    adisk = gmsh.model.occ.addDisk(xa, ya, 0, wire_radius, wire_radius)
    anode_wires.append(adisk)

    # Place IC wires only if flag is True
    if include_ic_wires:
        i1disk = gmsh.model.occ.addDisk(
            xloc_i1[i], yloc_i1[i], 0, wire_radius, wire_radius
        )
        i2disk = gmsh.model.occ.addDisk(
            xloc_i2[i], yloc_i2[i], 0, wire_radius, wire_radius
        )
        ic1_wires.append(i1disk)
        ic2_wires.append(i2disk)

anasen_barrel = gmsh.model.occ.addDisk(0, 0, 0, 500, 500)
# gmsh.model.occ.synchronize()
# gmsh.model.mesh.embed(1,anode_wires+cathode_wires,2,anasen_barrel)

gmsh.option.setNumber("Geometry.Tolerance", 1e-6)
gmsh.option.setNumber("Geometry.OCCFixDegenerated", 1)
gmsh.model.occ.synchronize()


# --- Surface Extraction ---
def get_surfs(disks):
    surfs = []
    for d in disks:
        surfs += [
            s[1] for s in gmsh.model.getBoundary([(2, d)], oriented=False) if s[0] == 1
        ]
    return surfs


if not 0 <= selected_cathode < len(cathode_wires):
    print("selected_cathode must be 0..%d" % (len(cathode_wires) - 1))
    quit()

cathode_wires_other = [
    disk for i, disk in enumerate(cathode_wires) if i != selected_cathode
]

# Cut the wire disks into the gas disk, replacing the old mesh.embed() call.
# embed only makes the mesher honour the curves; fragment splits the barrel
# into conforming surfaces, which is what lets one wire carry its own BC.
gmsh.model.occ.synchronize()

all_wire_disks = needle + guard_wires + cathode_wires + anode_wires
if include_ic_wires:
    all_wire_disks += ic1_wires + ic2_wires

gmsh.model.occ.fragment([(2, anasen_barrel)], [(2, d) for d in all_wire_disks])
gmsh.model.occ.synchronize()

# fragment re-derives the bounding curves, so re-extract them.
needle_surfs = get_surfs(needle) if include_needle else []
gwire_surfs = get_surfs(guard_wires)
awire_surfs = get_surfs(anode_wires)
cwire_surfs = get_surfs(cathode_wires_other)
cwire_sel_surfs = get_surfs([cathode_wires[selected_cathode]])
i1wire_surfs = get_surfs(ic1_wires) if include_ic_wires else []
i2wire_surfs = get_surfs(ic2_wires) if include_ic_wires else []

all_active_wire_surfs = (
    needle_surfs
    + gwire_surfs
    + awire_surfs
    + cwire_surfs
    + cwire_sel_surfs
    + i1wire_surfs
    + i2wire_surfs
)

f1 = gmsh.model.mesh.field.add("Distance")
gmsh.model.mesh.field.setNumbers(f1, "CurvesList", all_active_wire_surfs)

f2 = gmsh.model.mesh.field.add("Threshold")
gmsh.model.mesh.field.setNumber(f2, "InField", f1)
gmsh.model.mesh.field.setNumber(f2, "SizeMin", 0.05)  # Fine mesh near wires
gmsh.model.mesh.field.setNumber(f2, "SizeMax", 5.0)  # Large mesh in empty space
gmsh.model.mesh.field.setNumber(f2, "DistMin", 0.5)  # Apply SizeMin within 1mm
gmsh.model.mesh.field.setNumber(f2, "DistMax", 15.0)  # Transition to SizeMax by 20mm

gmsh.model.mesh.field.setAsBackgroundMesh(f2)


# --- Physical Groups ---
# These tags are the Target Bodies / Target Boundaries numbers in the .sif
# files. Do not renumber them without changing both sifs -- and note that
# ElmerGrid's -autoclean renumbers them for you, which silently breaks every
# boundary condition.
# Needle
if include_needle:
    gmsh.model.addPhysicalGroup(1, needle_surfs, tag=1, name="hot_needle")

# IC Wires
if include_ic_wires:
    gmsh.model.addPhysicalGroup(1, i1wire_surfs, tag=2, name="ic_wire_1")
    gmsh.model.addPhysicalGroup(1, i2wire_surfs, tag=3, name="ic_wire_2")

# Proportional Counter Wires
gmsh.model.addPhysicalGroup(1, gwire_surfs, tag=10, name="guard_wires")
gmsh.model.addPhysicalGroup(1, awire_surfs, tag=20, name="anode_wires")
gmsh.model.addPhysicalGroup(1, cwire_surfs, tag=30, name="cathode_wires")

# The one cathode the weighting potential is solved for, kept out of tag 30 so
# wires2d_weight.sif can drive it independently.
gmsh.model.addPhysicalGroup(1, cwire_sel_surfs, tag=40, name="cathode_wire_selected")

# Gas Volume (2D). fragment split the barrel into many surfaces, so collect
# them all; naming the original disk alone would solve one sliver of the domain.
all_surfaces_2d = [s[1] for s in gmsh.model.getEntities(dim=2)]
gmsh.model.addPhysicalGroup(2, all_surfaces_2d, tag=13, name="gas")

gmsh.option.setNumber("Mesh.Algorithm", 5)

gmsh.model.mesh.generate(2)
gmsh.model.mesh.setOrder(2)
# recombine() ran for over half an hour on the barrel surface without
# finishing, and refine() took the mesh to ~16M nodes. Neither is needed: the
# Threshold field above already gives 0.05 mm elements on 0.254 mm wires.
# gmsh.model.mesh.recombine()
# gmsh.model.mesh.refine()
gmsh.write("wires2d.msh")
# gmsh.fltk.run()
gmsh.finalize()