import os
import sys
import glob
import gc

# Limit the thread count for parallel solvers to prevent WSL from resource starvation/disconnecting
os.environ["OMP_NUM_THREADS"] = "6"
os.environ["MKL_NUM_THREADS"] = "6"
os.environ["TBB_NUM_THREADS"] = "6"
os.environ["OPENBLAS_NUM_THREADS"] = "6"
# Disable MKL internal memory manager to prevent memory accumulation (OOM) across solve calls
os.environ["MKL_DISABLE_FAST_MM"] = "1"

# Auto-detect and load Intel MKL before devsim is imported
mkl_libs = glob.glob(os.path.join(os.path.dirname(sys.executable), "../lib/libmkl_rt.so.*"))
if mkl_libs:
    os.environ["DEVSIM_MATH_LIBS"] = mkl_libs[0]

import devsim
import numpy as np

OUT_DIR = "output_this_run/"
os.makedirs(OUT_DIR, exist_ok=True)
import matplotlib.pyplot as plt
import time

# Enable Intel MKL PARDISO multi-threaded sparse solver
devsim.set_parameter(name="solver_type", value="pardiso")

sys.path.append("/home/pchan/devsim2026")
from device_config import *
from physics.model_create import *
from physics.new_physics import *

device = "device_2d"

# 1. Load the mesh
print("Loading mesh: device_2d.msh...")
devsim.create_gmsh_mesh(mesh=device, file="device_2d.msh")
devsim.add_gmsh_region(mesh=device, gmsh_name="Silicon", region="Silicon", material="Silicon")
devsim.add_gmsh_region(mesh=device, gmsh_name="Oxide", region="Oxide", material="Oxide")
devsim.add_gmsh_region(mesh=device, gmsh_name="Molding", region="Molding", material="Molding")

# Add contacts for Silicon region
devsim.add_gmsh_contact(mesh=device, gmsh_name="MT1_Si", name="MT1_Si", region="Silicon", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MT2_Si", name="MT2_Si", region="Silicon", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MT1_P12_Si", name="MT1_P12_Si", region="Silicon", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MT2_P12_Si", name="MT2_P12_Si", region="Silicon", material="metal")

# Add contacts for Oxide region
devsim.add_gmsh_contact(mesh=device, gmsh_name="MT1_Ox", name="MT1_Ox", region="Oxide", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MT2_Ox", name="MT2_Ox", region="Oxide", material="metal")

# Add contacts for Molding region
devsim.add_gmsh_contact(mesh=device, gmsh_name="MT1_Mold", name="MT1_Mold", region="Molding", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MT2_Mold", name="MT2_Mold", region="Molding", material="metal")

# Add interfaces
devsim.add_gmsh_interface(mesh=device, gmsh_name="Si_Ox_Interface", name="Si_Ox", region0="Silicon", region1="Oxide")
devsim.add_gmsh_interface(mesh=device, gmsh_name="Ox_Mold_Interface", name="Ox_Mold", region0="Oxide", region1="Molding")
devsim.add_gmsh_interface(mesh=device, gmsh_name="Si_Mold_Interface", name="Si_Mold", region0="Silicon", region1="Molding")

devsim.finalize_mesh(mesh=device)
devsim.create_device(mesh=device, device=device)

# 2. Set up doping in Silicon region
devsim.node_model(device=device, region="Silicon", name="nD_sub", equation=f"{N_SUB}")

def get_erfc_expr(peak, x1, x2, hdiff, vdiff):
    return f"{peak} * erfc(y / {vdiff}) * 0.5 * (erf((x - ({x1})) / {hdiff}) - erf((x - ({x2})) / {hdiff}))"

p11_left_expr = get_erfc_expr(P11_PEAK, -P11_X2, -P11_X1, P_WELL_HDDIFF, P_WELL_VDDIFF)
p11_right_expr = get_erfc_expr(P11_PEAK, P11_X1, P11_X2, P_WELL_HDDIFF, P_WELL_VDDIFF)
devsim.node_model(device=device, region="Silicon", name="nA_p11_l", equation=p11_left_expr)
devsim.node_model(device=device, region="Silicon", name="nA_p11_r", equation=p11_right_expr)

p12_left_expr = get_erfc_expr(P12_PEAK, -P12_X2, -P12_X1, P_WELL_HDDIFF, P_WELL_VDDIFF)
p12_right_expr = get_erfc_expr(P12_PEAK, P12_X1, P12_X2, P_WELL_HDDIFF, P_WELL_VDDIFF)
devsim.node_model(device=device, region="Silicon", name="nA_p12_l", equation=p12_left_expr)
devsim.node_model(device=device, region="Silicon", name="nA_p12_r", equation=p12_right_expr)

p13_left_expr = get_erfc_expr(P13_PEAK, -P13_X2, -P13_X1, P_WELL_HDDIFF, P_WELL_VDDIFF)
p13_right_expr = get_erfc_expr(P13_PEAK, P13_X1, P13_X2, P_WELL_HDDIFF, P_WELL_VDDIFF)
devsim.node_model(device=device, region="Silicon", name="nA_p13_l", equation=p13_left_expr)
devsim.node_model(device=device, region="Silicon", name="nA_p13_r", equation=p13_right_expr)

nplus_left_expr = get_erfc_expr(NPLUS_PEAK, -NPLUS_X2, -NPLUS_X1, NPLUS_HDDIFF, NPLUS_VDDIFF)
nplus_right_expr = get_erfc_expr(NPLUS_PEAK, NPLUS_X1, NPLUS_X2, NPLUS_HDDIFF, NPLUS_VDDIFF)
devsim.node_model(device=device, region="Silicon", name="nD_nplus_l", equation=nplus_left_expr)
devsim.node_model(device=device, region="Silicon", name="nD_nplus_r", equation=nplus_right_expr)

mring_l_expr = get_erfc_expr(NPLUS_PEAK, -W_DEVICE, -MRING_X1, NPLUS_HDDIFF, NPLUS_VDDIFF)
mring_r_expr = get_erfc_expr(NPLUS_PEAK, MRING_X1, W_DEVICE, NPLUS_HDDIFF, NPLUS_VDDIFF)
devsim.node_model(device=device, region="Silicon", name="nD_mring_l", equation=mring_l_expr)
devsim.node_model(device=device, region="Silicon", name="nD_mring_r", equation=mring_r_expr)

devsim.node_model(device=device, region="Silicon", name="Donors", 
                  equation="nD_sub + nD_nplus_l + nD_nplus_r + nD_mring_l + nD_mring_r")
devsim.node_model(device=device, region="Silicon", name="Acceptors", 
                  equation="1e10 + nA_p11_l + nA_p11_r + nA_p12_l + nA_p12_r + nA_p13_l + nA_p13_r")
devsim.node_model(device=device, region="Silicon", name="NetDoping", equation="Donors - Acceptors")
devsim.node_model(device=device, region="Silicon", name="LogNetDoping", equation="asinh(NetDoping / 2.0) / log(10.0)")

# 3. Initialize electrostatic potential simulation (Poisson only)
CreateSolution(device, "Silicon", "Potential")
devsim.set_parameter(device=device, name="T", value="300")
CreateSiliconPotentialOnly(device, "Silicon")

# Oxide potential equations
def CreateOxidePotentialOnly(device, region):
    if not InNodeModelList(device, region, "Potential"):
        CreateSolution(device, region, "Potential")
    devsim.set_parameter(device=device, region=region, name="Permittivity", value=3.9 * 8.85e-14)
    efield = "(Potential@n0 - Potential@n1)*EdgeInverseLength"
    CreateEdgeModel(device, region, "EField", efield)
    CreateEdgeModelDerivatives(device, region, "EField", efield, "Potential")
    dfield = "Permittivity*EField"
    CreateEdgeModel(device, region, "PotentialEdgeFlux", dfield)
    CreateEdgeModelDerivatives(device, region, "PotentialEdgeFlux", dfield, "Potential")
    devsim.equation(device=device, region=region, name="PotentialEquation", variable_name="Potential",
                    edge_model="PotentialEdgeFlux", variable_update="default", min_error=1e-3)

CreateOxidePotentialOnly(device, "Oxide")

# Molding potential equations
def CreateMoldingPotentialOnly(device, region):
    if not InNodeModelList(device, region, "Potential"):
        CreateSolution(device, region, "Potential")
    devsim.set_parameter(device=device, region=region, name="Permittivity", value=4.0 * 8.85e-14)
    efield = "(Potential@n0 - Potential@n1)*EdgeInverseLength"
    CreateEdgeModel(device, region, "EField", efield)
    CreateEdgeModelDerivatives(device, region, "EField", efield, "Potential")
    dfield = "Permittivity*EField"
    CreateEdgeModel(device, region, "PotentialEdgeFlux", dfield)
    CreateEdgeModelDerivatives(device, region, "PotentialEdgeFlux", dfield, "Potential")
    devsim.equation(device=device, region=region, name="PotentialEquation", variable_name="Potential",
                    edge_model="PotentialEdgeFlux", variable_update="default", min_error=1e-3)

CreateMoldingPotentialOnly(device, "Molding")

# Interfaces continuous potential
def CreateContinuousPotentialInterface(device, interface):
    model_name = CreateContinuousInterfaceModel(device, interface, "Potential")
    devsim.interface_equation(device=device, interface=interface, name="PotentialEquation",
                              interface_model=model_name, type="continuous")

CreateContinuousPotentialInterface(device, "Si_Ox")
CreateContinuousPotentialInterface(device, "Ox_Mold")
CreateContinuousPotentialInterface(device, "Si_Mold")

# Potential contacts setup
silicon_contacts = ["MT1_Si", "MT2_Si", "MT1_P12_Si", "MT2_P12_Si"]
for c in silicon_contacts:
    devsim.set_parameter(device=device, name=GetContactBiasName(c), value=0.0)
    CreateSiliconPotentialOnlyContact(device, "Silicon", c)

def CreateOxidePotentialOnlyContact(device, region, contact):
    contact_bias = GetContactBiasName(contact)
    contact_model = f"Potential - {contact_bias}"
    contact_model_name = f"{contact}_bc"
    CreateContactNodeModel(device, contact, contact_model_name, contact_model)
    CreateContactNodeModelDerivative(device, contact, contact_model_name, contact_model, "Potential")
    devsim.contact_equation(device=device, contact=contact, name="PotentialEquation",
                            node_model=contact_model_name, edge_charge_model="PotentialEdgeFlux")

oxide_contacts = ["MT1_Ox", "MT2_Ox"]
for c in oxide_contacts:
    devsim.set_parameter(device=device, name=GetContactBiasName(c), value=0.0)
    CreateOxidePotentialOnlyContact(device, "Oxide", c)

def CreateMoldingPotentialOnlyContact(device, region, contact):
    contact_bias = GetContactBiasName(contact)
    contact_model = f"Potential - {contact_bias}"
    contact_model_name = f"{contact}_bc"
    CreateContactNodeModel(device, contact, contact_model_name, contact_model)
    CreateContactNodeModelDerivative(device, contact, contact_model_name, contact_model, "Potential")
    devsim.contact_equation(device=device, contact=contact, name="PotentialEquation",
                            node_model=contact_model_name, edge_charge_model="PotentialEdgeFlux")

molding_contacts = ["MT1_Mold", "MT2_Mold"]
for c in molding_contacts:
    devsim.set_parameter(device=device, name=GetContactBiasName(c), value=0.0)
    CreateMoldingPotentialOnlyContact(device, "Molding", c)

# Solve initial zero-bias Poisson
print("Solving initial Poisson at thermal equilibrium...")
devsim.solve(type="dc", absolute_error=1.0, relative_error=1e-10, maximum_iterations=50)
print("Initial Poisson converged.")

# 4. Set up carrier solutions for Silicon Drift-Diffusion
# Compute initial guess for Electrons and Holes based on Potential
CreateSolution(device, "Silicon", "Electrons")
CreateSolution(device, "Silicon", "Holes")

devsim.set_node_values(device=device, region="Silicon", name="Electrons", init_from="IntrinsicElectrons")
devsim.set_node_values(device=device, region="Silicon", name="Holes", init_from="IntrinsicHoles")

# Redefine IntrinsicElectrons, IntrinsicHoles, and related models to avoid potential exponential overflow at high bias.
print("Redefining equilibrium models to prevent high-bias exponential overflow...")
devsim.node_model(device=device, region="Silicon", name="IntrinsicElectrons", equation="Electrons")
devsim.node_model(device=device, region="Silicon", name="IntrinsicElectrons:Potential", equation="0")
devsim.node_model(device=device, region="Silicon", name="IntrinsicElectrons:Electrons", equation="1")
devsim.node_model(device=device, region="Silicon", name="IntrinsicElectrons:Holes", equation="0")

devsim.node_model(device=device, region="Silicon", name="IntrinsicHoles", equation="Holes")
devsim.node_model(device=device, region="Silicon", name="IntrinsicHoles:Potential", equation="0")
devsim.node_model(device=device, region="Silicon", name="IntrinsicHoles:Electrons", equation="0")
devsim.node_model(device=device, region="Silicon", name="IntrinsicHoles:Holes", equation="1")

devsim.node_model(device=device, region="Silicon", name="IntrinsicCharge", equation="Holes - Electrons + NetDoping")
devsim.node_model(device=device, region="Silicon", name="IntrinsicCharge:Potential", equation="0")
devsim.node_model(device=device, region="Silicon", name="IntrinsicCharge:Electrons", equation="-1")
devsim.node_model(device=device, region="Silicon", name="IntrinsicCharge:Holes", equation="1")

devsim.node_model(device=device, region="Silicon", name="PotentialIntrinsicCharge", equation="0")
devsim.node_model(device=device, region="Silicon", name="PotentialIntrinsicCharge:Potential", equation="0")

# Mobility and drift diffusion equations
opts = CreateAroraMobilityLF(device, "Silicon")
# Bypassing HFMobility to prevent zero-bias convergence oscillations
CreateSiliconDriftDiffusion(device, "Silicon", **opts)
devsim.node_model(device=device, region="Silicon", name="LogElectrons", equation="log(Electrons + 1e-10) / log(10.0)")
devsim.node_model(device=device, region="Silicon", name="LogHoles", equation="log(Holes + 1e-10) / log(10.0)")

# Re-setup Silicon contacts for Drift-Diffusion
for c in silicon_contacts:
    CreateSiliconDriftDiffusionContact(device, "Silicon", c, opts['Jn'], opts['Jp'])

# Solve initial zero-bias Drift-Diffusion with standard tolerances (using default log_damp updates)
print("Solving initial Drift-Diffusion equations at zero bias...")
devsim.solve(type="dc", absolute_error=1e10, relative_error=1e30, charge_error=1e12, maximum_iterations=50)
print("Initial Drift-Diffusion converged successfully!")

# Switch continuity and potential equations for the bias sweep
print("Configuring continuity and potential equations for the bias sweep (min_error=1e5, positive update)...")

# Instantiate Avalanche (Impact Ionization) edge generation model
CreateAvalancheGeneration(device, "Silicon", opts['Jn'], opts['Jp'])

devsim.equation(device=device, region="Silicon", name="ElectronContinuityEquation", variable_name="Electrons",
                time_node_model="NCharge", edge_model=opts['Jn'], edge_volume_model="", 
                variable_update="positive", node_model="ElectronGeneration", min_error=1e5)
devsim.equation(device=device, region="Silicon", name="HoleContinuityEquation", variable_name="Holes",
                time_node_model="PCharge", edge_model=opts['Jp'], edge_volume_model="", 
                variable_update="positive", node_model="HoleGeneration", min_error=1e5)
devsim.equation(device=device, region="Silicon", name="PotentialEquation", variable_name="Potential",
                node_model="PotentialNodeCharge", edge_model="DField", variable_update="default", min_error=1e-3)

# Save zero-bias tecplot and VTK
devsim.write_devices(file=f"{OUT_DIR}sweep_preview_0V.tec", type="tecplot")
# devsim.write_devices(file="sweep_preview_0V", type="vtk")

# 5. Define Sweep Parameters
v_target = 1000.0
v_current = 0.0
step_size = 0.1         # Initial step size (V)
max_step = 50.0         # Maximum step size (V)
min_step = 1e-4         # Minimum step size (V)
compliance_current = 1e-3  # 1 mA compliance current

# Helper functions to save/restore state in case of convergence failure
def save_state(device):
    state = {}
    for region in ["Silicon", "Oxide", "Molding"]:
        state[region] = {
            "Potential": list(devsim.get_node_model_values(device=device, region=region, name="Potential"))
        }
    state["Silicon"]["Electrons"] = list(devsim.get_node_model_values(device=device, region="Silicon", name="Electrons"))
    state["Silicon"]["Holes"] = list(devsim.get_node_model_values(device=device, region="Silicon", name="Holes"))
    return state

def restore_state(device, state):
    for region in ["Silicon", "Oxide", "Molding"]:
        devsim.set_node_values(device=device, region=region, name="Potential", values=state[region]["Potential"])
    devsim.set_node_values(device=device, region="Silicon", name="Electrons", values=state["Silicon"]["Electrons"])
    devsim.set_node_values(device=device, region="Silicon", name="Holes", values=state["Silicon"]["Holes"])

# File logging setup
time_log = open(f"{OUT_DIR}simulation_time.log", "w", buffering=1)
time_log.write("Time\tVoltage(V)\tStep(V)\tCurrent(A)\tIterations\tTimeTaken(s)\n")

# Arrays to store I-V data
voltage_list = [0.0]
current_list = [0.0]

# Save initial state
state = save_state(device)
start_sweep_time = time.time()

print("Beginning adaptive bias sweep...")
step_count = 0

# Targets for saving intermediate state checkpoints
save_targets = [5.0, 50.0, 500.0]
saved_targets = set()

while v_current < v_target:
    v_next = min(v_current + step_size, v_target)
    
    # Apply new bias values to MT1 contacts
    for c in ["MT1_Si", "MT1_P12_Si", "MT1_Ox", "MT1_Mold"]:
        devsim.set_parameter(device=device, name=f"{c}_bias", value=v_next)
        
    step_start_time = time.time()
    try:
        # Stage 1: Pre-conditioning (Relaxed Tolerance to get a good initial guess)
        iters1 = 10
        try:
            res1 = devsim.solve(type="dc", absolute_error=1e10, relative_error=1e-1, charge_error=1e12, maximum_iterations=10, info=True)
            iters1 = len(res1.get("iterations", []))
        except devsim.error:
            pass # Ignore non-convergence in pre-conditioning
        
        import psutil
        mem_stage1 = psutil.Process(os.getpid()).memory_info().rss / (1024**3)
        print(f"Stage 1 (Precondition) Memory Usage: {mem_stage1:.1f} GB")
            
        # Stage 2: Precision Newton (Strict Tolerance) - slightly relaxed relative_error to 3e-3 to avoid limit cycle oscillations
        res = devsim.solve(type="dc", absolute_error=1e10, relative_error=3e-3, charge_error=1e12, maximum_iterations=12, info=True)
        iters2 = len(res.get("iterations", []))
        
        mem_stage2 = psutil.Process(os.getpid()).memory_info().rss / (1024**3)
        print(f"Stage 2 (Newton) Memory Usage: {mem_stage2:.1f} GB")
        iters = f"{iters1}+{iters2}"
        total_iters = iters1 + iters2
        
        if not res.get("converged", False):
            raise devsim.error("Convergence failure")
            
        step_end_time = time.time()
        time_taken = step_end_time - step_start_time
        
        # Convergence succeeded! Compute current at MT1 terminal
        # MT1 terminal current is the sum of currents on MT1_Si and MT1_P12_Si
        i_n_si = devsim.get_contact_current(device=device, contact="MT1_Si", equation="ElectronContinuityEquation")
        i_p_si = devsim.get_contact_current(device=device, contact="MT1_Si", equation="HoleContinuityEquation")
        i_n_p12 = devsim.get_contact_current(device=device, contact="MT1_P12_Si", equation="ElectronContinuityEquation")
        i_p_p12 = devsim.get_contact_current(device=device, contact="MT1_P12_Si", equation="HoleContinuityEquation")
        
        total_curr = i_n_si + i_p_si + i_n_p12 + i_p_p12
        
        # Update simulation status
        v_current = v_next
        state = save_state(device)
        
        voltage_list.append(v_current)
        current_list.append(total_curr)
        
        print(f"Step {step_count}: Converged at V = {v_current:.4f} V, I = {total_curr:.4e} A. Step size: {step_size:.4f} V. Iterations: {iters}. Time: {time_taken:.2f} s")
        
        # Log to file
        time_log.write(f"{time.strftime('%X')}\t{v_current:.4f}\t{step_size:.4f}\t{total_curr:.4e}\t{iters}\t{time_taken:.2f}\t{mem_stage1:.1f}GB\t{mem_stage2:.1f}GB\n")
        
        # Save checkpoints when crossing target voltages
        for target in save_targets:
            if v_current >= target and target not in saved_targets:
                filename = f"sweep_preview_{int(target)}V.tec"
                filename_vtk = f"sweep_preview_{int(target)}V"
                print(f"Saving checkpoint at V = {v_current:.2f} V to {filename} and VTK...")
                devsim.write_devices(file=f"{OUT_DIR}{filename}", type="tecplot")
                # devsim.write_devices(file=filename_vtk, type="vtk")
                saved_targets.add(target)
        
        # Compliance check
        if abs(total_curr) >= compliance_current:
            print(f"Compliance current of {compliance_current:.1e} A reached at V = {v_current:.4f} V. Stopping sweep.")
            time_log.write(f"Compliance current reached at V = {v_current:.4f} V.\n")
            break
            
        # Grow step size for next step adaptively based on Newton iterations
        if total_iters <= 6:
            step_size = min(step_size * 1.2, max_step)
        elif total_iters <= 9:
            pass # Keep step size the same
        else:
            step_size = max(step_size * 0.8, min_step)
            
        step_count += 1
        
    except devsim.error as e:
        # Convergence failure: restore last state and cut step size
        step_end_time = time.time()
        time_taken = step_end_time - step_start_time
        print(f"Convergence failure at V = {v_next:.4f} V. Restoring state and scaling step size by 0.577 from {step_size:.4f} V.")
        time_log.write(f"{time.strftime('%X')}\t{v_next:.4f}\t{step_size:.4f}\tFAILED\t-\t{time_taken:.2f}\n")
        
        restore_state(device, state)
        step_size *= 0.577
        
        if step_size < min_step:
            print("Step size has fallen below minimum limit. Aborting simulation.")
            time_log.write(f"Aborted: step size fell below {min_step:.1e} V\n")
            break

total_sweep_time = time.time() - start_sweep_time
print(f"Sweep completed in {total_sweep_time:.2f} s.")
time_log.write(f"Total Sweep Time: {total_sweep_time:.2f} s\n")
time_log.close()

# 6. Save final results and generate plots
# Save final tecplot and VTK at highest voltage
devsim.write_devices(file=f"{OUT_DIR}sweep_preview_final.tec", type="tecplot")
# devsim.write_devices(file="sweep_preview_final", type="vtk")

# Save I-V data to CSV
np.savetxt(f"{OUT_DIR}sweep_iv_2d.csv", np.column_stack((voltage_list, current_list)), 
           header="Voltage(V),Current(A)", delimiter=",")

# Plot and save I-V curve
plt.figure(figsize=(8, 6))
plt.plot(voltage_list, np.abs(current_list), 'o-', color='#1f77b4', markersize=4)
plt.yscale('log')
plt.grid(True, which="both", ls="--")
plt.xlabel("Bias Voltage (V)")
plt.ylabel("Terminal Current Magnitude (A)")
plt.title("TVS 2D Bidirectional Bias Sweep I-V Curve (Log Scale)")
plt.tight_layout()
plt.savefig(f"{OUT_DIR}sweep_iv_2d.png", dpi=300)
plt.close()

# Generate potential & electric field plots at final converged bias
# Extract final node values
x_si = np.array(devsim.get_node_model_values(device=device, region="Silicon", name="x")) / um
y_si = np.array(devsim.get_node_model_values(device=device, region="Silicon", name="y")) / um
pot_si = np.array(devsim.get_node_model_values(device=device, region="Silicon", name="Potential"))
tri_si = np.array(devsim.get_element_node_list(device=device, region="Silicon"))

# Compute final Emag
for reg in ["Silicon", "Oxide", "Molding"]:
    devsim.element_from_edge_model(edge_model="EField", device=device, region=reg)
    devsim.element_model(device=device, region=reg, name="Emag", equation="(EField_x^2 + EField_y^2)^(0.5)")
    
emag_si = np.array(devsim.get_element_model_values(device=device, region="Silicon", name="Emag"))[::3]

x_ox = np.array(devsim.get_node_model_values(device=device, region="Oxide", name="x")) / um
y_ox = np.array(devsim.get_node_model_values(device=device, region="Oxide", name="y")) / um
pot_ox = np.array(devsim.get_node_model_values(device=device, region="Oxide", name="Potential"))
tri_ox = np.array(devsim.get_element_node_list(device=device, region="Oxide"))
emag_ox = np.array(devsim.get_element_model_values(device=device, region="Oxide", name="Emag"))[::3]

x_mold = np.array(devsim.get_node_model_values(device=device, region="Molding", name="x")) / um
y_mold = np.array(devsim.get_node_model_values(device=device, region="Molding", name="y")) / um
pot_mold = np.array(devsim.get_node_model_values(device=device, region="Molding", name="Potential"))
tri_mold = np.array(devsim.get_element_node_list(device=device, region="Molding"))
emag_mold = np.array(devsim.get_element_model_values(device=device, region="Molding", name="Emag"))[::3]

def draw_device_boundaries(ax):
    ax.plot([-W_DEVICE/um, W_DEVICE/um], [-T_OX/um, -T_OX/um], color='black', linestyle='--', linewidth=0.8)
    ax.plot([-W_DEVICE/um, W_DEVICE/um], [0, 0], color='black', linestyle='-', linewidth=0.8)
    ax.plot([-W_DEVICE/um, -W_DEVICE/um], [0, H_SI/um], color='black', linestyle='-', linewidth=0.8)
    ax.plot([W_DEVICE/um, W_DEVICE/um], [0, H_SI/um], color='black', linestyle='-', linewidth=0.8)
    ax.plot([-W_SIM/um, W_SIM/um], [H_SI/um, H_SI/um], color='black', linestyle='-', linewidth=1.2)

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 14))

# Plot Potential
tcf1_si = ax1.tripcolor(x_si, y_si, tri_si, pot_si, cmap='RdYlBu_r', shading='gouraud')
tcf1_ox = ax1.tripcolor(x_ox, y_ox, tri_ox, pot_ox, cmap='RdYlBu_r', shading='gouraud')
tcf1_mold = ax1.tripcolor(x_mold, y_mold, tri_mold, pot_mold, cmap='RdYlBu_r', shading='gouraud')
fig.colorbar(tcf1_si, ax=ax1, label='Electrostatic Potential (V)')
draw_device_boundaries(ax1)
ax1.set_xlabel('X (μm)')
ax1.set_ylabel('Y (μm)')
ax1.set_title(f'2D Electrostatic Potential at V = {v_current:.2f} V')
ax1.set_xlim(-W_SIM / um, W_SIM / um)
ax1.set_ylim(H_SI/um + 15.0, -110.0)

# Plot EField Magnitude
tcf2_si = ax2.tripcolor(x_si, y_si, tri_si, facecolors=emag_si, cmap='inferno', shading='flat')
tcf2_ox = ax2.tripcolor(x_ox, y_ox, tri_ox, facecolors=emag_ox, cmap='inferno', shading='flat')
tcf2_mold = ax2.tripcolor(x_mold, y_mold, tri_mold, facecolors=emag_mold, cmap='inferno', shading='flat')
fig.colorbar(tcf2_si, ax=ax2, label='Electric Field Magnitude (V/cm)')
draw_device_boundaries(ax2)
ax2.set_xlabel('X (μm)')
ax2.set_ylabel('Y (μm)')
ax2.set_title(f'2D Electric Field Magnitude at V = {v_current:.2f} V')
ax2.set_xlim(-W_SIM / um, W_SIM / um)
ax2.set_ylim(H_SI/um + 15.0, -110.0)

plt.tight_layout()
plt.savefig(f"{OUT_DIR}sweep_potential_2d.png", dpi=300)
plt.close()

print(f"Sweep visualization plots saved: sweep_iv_2d.png and sweep_potential_2d.png.")
import os; os.system(f'{sys.executable} plot_speed.py')