import os
import sys

# Limit the thread count for parallel solvers to prevent WSL from resource starvation/disconnecting
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_NUM_THREADS"] = "4"
os.environ["TBB_NUM_THREADS"] = "4"
os.environ["OPENBLAS_NUM_THREADS"] = "4"

import devsim
import numpy as np

DEV_DIR = os.environ.get("DEV_DIR", "devices/Triac_rp")
sys.path.insert(0, os.path.abspath(DEV_DIR))

OUT_DIR = os.path.join(os.environ.get("OUT_DIR", os.path.join(DEV_DIR, "output_this_run")), "")
os.makedirs(OUT_DIR, exist_ok=True)
import matplotlib.pyplot as plt
from device_config import *
from physics.model_create import *
from physics.new_physics import *

device = "device_2d"

# 1. Load the mesh
mesh_file = os.path.join(DEV_DIR, "device_2d.msh")
devsim.create_gmsh_mesh(mesh=device, file=mesh_file)
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 (MT1, MT2, and P12 virtual contacts; MRING and Substrate Bottom will float as Neumann boundaries)
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)
# --- rest of file ---
# Skip lines 35-124 as they are unchanged


# 2. Set up doping in Silicon region
if os.environ.get("USE_PCAD", "false").lower() == "true":
    from device_pcad_config import apply_pcad_doping_2d
    apply_pcad_doping_2d(device, region="Silicon")
else:
    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}))"

    # P-wells
    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)

    # N+
    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
    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)

    # Combine into Donors and Acceptors
    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. Create solution variables and physics models
CreateSolution(device, "Silicon", "Potential")
sim_temp = os.environ.get("TEMP", "300")
devsim.set_parameter(device=device, name="T", value=sim_temp)
CreateSiliconPotentialOnly(device, "Silicon")

# Oxide Potential physics setup
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")

CreateOxidePotentialOnly(device, "Oxide")

# Molding Potential physics setup
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")

CreateMoldingPotentialOnly(device, "Molding")

# Interfaces (continuous electrostatic 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")

# 4. Apply contacts boundary conditions
# Silicon contacts
silicon_contacts = ["MT1_Si", "MT2_Si"]
for c in silicon_contacts:
    devsim.set_parameter(device=device, name=GetContactBiasName(c), value=0.0)
    CreateSiliconPotentialOnlyContact(device, "Silicon", c)

# P12 Virtual Silicon contacts (tied to MT1 and MT2 respectively)
devsim.set_parameter(device=device, name="MT1_P12_Si_bias", value=0.0)
CreateSiliconPotentialOnlyContact(device, "Silicon", "MT1_P12_Si")

devsim.set_parameter(device=device, name="MT2_P12_Si_bias", value=0.0)
CreateSiliconPotentialOnlyContact(device, "Silicon", "MT2_P12_Si")

# Oxide contacts
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)

# Molding contacts
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)

# 5. Solve Potential at equilibrium (zero bias)
print("Solving Poisson/Laplace equations at thermal equilibrium...")
devsim.solve(type="dc", absolute_error=1.0, relative_error=1e-10, maximum_iterations=50)
print("Solution converged successfully!")

# Compute electric field magnitude (Emag) on elements
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)")

# Save the solution to static_preview.tec and static_preview.vtm
devsim.write_devices(file=f"{OUT_DIR}static_preview.tec", type="tecplot")
# devsim.write_devices(file="static_preview", type="vtk")
print("Saved static_preview.tec and static_preview.vtm (VTK) for ParaView.")

# 6. Extract data and generate a Matplotlib plot
print("Extracting data for plotting...")

# Silicon region data
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"))
emag_si = np.array(devsim.get_element_model_values(device=device, region="Silicon", name="Emag"))[::3]

# Oxide region data
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]

# Molding region data
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):
    # Overlay lines for regions
    # Oxide Top: Y = -T_OX from -W_DEVICE to W_DEVICE
    ax.plot([-W_DEVICE/um, W_DEVICE/um], [-T_OX/um, -T_OX/um], color='black', linestyle='--', linewidth=0.8)
    # Silicon-Oxide Interface: Y = 0 from -W_DEVICE to W_DEVICE
    ax.plot([-W_DEVICE/um, W_DEVICE/um], [0, 0], color='black', linestyle='-', linewidth=0.8)
    # Silicon Die Side Boundaries: X = +-W_DEVICE from Y = 0 to H_SI
    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)
    # Bottom: Y = H_SI from -W_SIM to W_SIM
    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('2D Electrostatic Potential at Zero Bias (Floating Bottom & MRING)')
ax1.set_xlim(-W_SIM / um, W_SIM / um)
ax1.set_ylim(H_SI/um + 15.0, -110.0)

# Plot Electric Field Magnitude (Emag)
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('2D Electric Field Magnitude at Zero Bias (Floating Bottom & MRING)')
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}static_potential_2d.png", dpi=300)
plt.close()
print("Plot saved to static_potential_2d.png")