import devsim
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.tri as tri
from device_config import *

device = "device_2d"

# 1. Load the mesh
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="MRING_L_Si", name="MRING_L", region="Silicon", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MRING_R_Si", name="MRING_R", region="Silicon", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="Substrate_Bottom", name="Substrate_Bottom", 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")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MRING_L_Ox", name="MRING_L_Ox", region="Oxide", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MRING_R_Ox", name="MRING_R_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")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MRING_L_Mold", name="MRING_L_Mold", region="Molding", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="MRING_R_Mold", name="MRING_R_Mold", region="Molding", material="metal")
devsim.add_gmsh_contact(mesh=device, gmsh_name="Substrate_Bottom_Mold", name="Substrate_Bottom_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. Define Doping Profiles using sub-models to avoid long strings
# Substrate (N-type)
devsim.node_model(device=device, region="Silicon", name="nD_sub", equation=f"{N_SUB}")

# Helper to generate 2D erfc profile string
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-well profiles (p11, p12, p13 on both sides)
# p11
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
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
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+ profiles
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 N+ profiles
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")

# NetDoping
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)")
devsim.node_model(device=device, region="Silicon", name="LogAcceptors", equation="log(Acceptors) / log(10.0)")

# Write Tecplot output for ParaView
devsim.write_devices(file="device_2d.tec", type="tecplot")
devsim.write_devices(file="preview.tec", type="tecplot")
print("Saved device_2d.tec and preview.tec")

# 4. Generate a 2D Plot with Matplotlib to verify the doping profile
print("Generating 2D plot...")
x = devsim.get_node_model_values(device=device, region="Silicon", name="x")
y = devsim.get_node_model_values(device=device, region="Silicon", name="y")
net_dop = devsim.get_node_model_values(device=device, region="Silicon", name="NetDoping")
log_dop = devsim.get_node_model_values(device=device, region="Silicon", name="LogNetDoping")

elements = devsim.get_element_node_list(device=device, region="Silicon")

# Convert elements into a format usable by matplotlib
triangles = np.array(elements)

# scale to micrometers for plotting
x_um = np.array(x) / um
y_um = np.array(y) / um

log_acceptors = devsim.get_node_model_values(device=device, region="Silicon", name="LogAcceptors")

def draw_oxide_and_metal(ax):
    # 0. Molding Region (light yellow-gray or beige)
    # Top Molding
    ax.add_patch(plt.Rectangle((-W_SIM / um, - (T_OX + H_MOLD) / um), 2 * W_SIM / um, H_MOLD / um, facecolor='#fbfcf7', edgecolor='lightgray', linewidth=0.5, alpha=0.9))
    # Left Side Molding
    ax.add_patch(plt.Rectangle((-W_SIM / um, - T_OX / um), (W_SIM - W_DEVICE) / um, (H_SI + T_OX) / um, facecolor='#fbfcf7', edgecolor='lightgray', linewidth=0.5, alpha=0.9))
    # Right Side Molding
    ax.add_patch(plt.Rectangle((W_DEVICE / um, - T_OX / um), (W_SIM - W_DEVICE) / um, (H_SI + T_OX) / um, facecolor='#fbfcf7', edgecolor='lightgray', linewidth=0.5, alpha=0.9))
    
    # 1. Oxide Layer (light blue-gray)
    rect_oxide = plt.Rectangle((-W_DEVICE / um, - T_OX / um), 2 * W_DEVICE / um, T_OX / um, facecolor='#eaeef2', edgecolor='gray', linewidth=0.5, alpha=0.9)
    ax.add_patch(rect_oxide)
    
    # 2. Leadframe Island at bottom (dark grey-blue)
    rect_leadframe = plt.Rectangle((-W_SIM / um, H_SI / um), 2 * W_SIM / um, 15.0, facecolor='#34495e', edgecolor='black', alpha=0.9)
    ax.add_patch(rect_leadframe)
    
    # 3. Metal color & settings (grey color for electrodes)
    m_color = '#7f8c8d'  # sleek dark gray
    m_edge = '#2c3e50'
    m_alpha = 1.0
    
    # MT1 (Right side)
    # Long plate top part: X in [30, 186], Y in [-2.5, -2.0]
    ax.add_patch(plt.Rectangle((30.0, -2.5), 156.0, 0.5, facecolor=m_color, edgecolor=m_edge, alpha=m_alpha))
    # Via 1 (under p11): X in [82.5, 92.5], Y in [-2.0, 0]
    ax.add_patch(plt.Rectangle((82.5, -2.0), 10.0, 2.0, facecolor=m_color, edgecolor=m_edge, alpha=m_alpha))
    # Via 2 (under p13): X in [169.5, 179.5], Y in [-2.0, 0]
    ax.add_patch(plt.Rectangle((169.5, -2.0), 10.0, 2.0, facecolor=m_color, edgecolor=m_edge, alpha=m_alpha))
    # Small plate top part: X in [250, 295], Y in [-2.5, -2.0]
    ax.add_patch(plt.Rectangle((250.0, -2.5), 45.0, 0.5, facecolor=m_color, edgecolor=m_edge, alpha=m_alpha))
    
    # MT2 (Left side)
    # Long plate top part: X in [-186, -30], Y in [-2.5, -2.0]
    ax.add_patch(plt.Rectangle((-186.0, -2.5), 156.0, 0.5, facecolor=m_color, edgecolor=m_edge, alpha=m_alpha))
    # Via 1: X in [-92.5, -82.5], Y in [-2.0, 0]
    ax.add_patch(plt.Rectangle((-92.5, -2.0), 10.0, 2.0, facecolor=m_color, edgecolor=m_edge, alpha=m_alpha))
    # Via 2: X in [-179.5, -169.5], Y in [-2.0, 0]
    ax.add_patch(plt.Rectangle((-179.5, -2.0), 10.0, 2.0, facecolor=m_color, edgecolor=m_edge, alpha=m_alpha))
    # Small plate top part: X in [-295, -250], Y in [-2.5, -2.0]
    ax.add_patch(plt.Rectangle((-295.0, -2.5), 45.0, 0.5, facecolor=m_color, edgecolor=m_edge, alpha=m_alpha))
    
    # MRING (Right & Left)
    mring_color = '#e67e22'  # bright orange-red for MRING to distinguish
    mring_edge = '#d35400'
    # Right MRING via: X in [340, 356], Y in [-2.0, 0]
    ax.add_patch(plt.Rectangle((340.0, -2.0), 16.0, 2.0, facecolor=mring_color, edgecolor=mring_edge, alpha=m_alpha))
    # Right MRING top plate: X in [335, 356], Y in [-2.5, -2.0]
    ax.add_patch(plt.Rectangle((335.0, -2.5), 21.0, 0.5, facecolor=mring_color, edgecolor=mring_edge, alpha=m_alpha))
    
    # Left MRING via: X in [-356, -340], Y in [-2.0, 0]
    ax.add_patch(plt.Rectangle((-356.0, -2.0), 16.0, 2.0, facecolor=mring_color, edgecolor=mring_edge, alpha=m_alpha))
    # Left MRING top plate: X in [-356, -335], Y in [-2.5, -2.0]
    ax.add_patch(plt.Rectangle((-356.0, -2.5), 21.0, 0.5, facecolor=mring_color, edgecolor=mring_edge, alpha=m_alpha))

    # Add text labels
    ax.text(108.0, -2.8, 'MT1', color='black', fontsize=8, ha='center', weight='bold')
    ax.text(-108.0, -2.8, 'MT2', color='black', fontsize=8, ha='center', weight='bold')
    ax.text(348.0, -2.8, 'MRING', color='#d35400', fontsize=8, ha='center', weight='bold')
    ax.text(-348.0, -2.8, 'MRING', color='#d35400', fontsize=8, ha='center', weight='bold')
    ax.text(0, -50.0, 'Molding Region', color='darkgreen', fontsize=9, ha='center', va='center')
    ax.text(-406.0, 50.0, 'Molding\nCompound\n(Side)', color='darkgreen', fontsize=8, ha='center', va='center')
    ax.text(406.0, 50.0, 'Molding\nCompound\n(Side)', color='darkgreen', fontsize=8, ha='center', va='center')
    ax.text(0, -1.0, 'Oxide', color='blue', fontsize=9, ha='center', va='center')
    ax.text(0, H_SI/um + 7.5, 'Leadframe paddle (Island)', color='white', fontsize=9, ha='center', va='center', weight='bold')

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

# Subplot 1: Net Doping
tcf1 = ax1.tripcolor(x_um, y_um, triangles, log_dop, cmap='coolwarm', shading='flat')
fig.colorbar(tcf1, ax=ax1, label='Log10(NetDoping) [asinh(N/2)/log(10)]')
draw_oxide_and_metal(ax1)
ax1.set_xlabel('X (μm)')
ax1.set_ylabel('Y (μm)')
ax1.set_title('2D Net Doping Profile (NetDoping = Donors - Acceptors)')
ax1.set_xlim(-W_SIM / um, W_SIM / um)
ax1.set_ylim(H_SI/um + 15.0, -110.0) # Show substrate, bottom contact, oxide, and top molding

# Subplot 2: Acceptors (P-type dopants)
tcf2 = ax2.tripcolor(x_um, y_um, triangles, np.array(log_acceptors), cmap='Purples', shading='flat')
fig.colorbar(tcf2, ax=ax2, label='Log10(Acceptor Doping)')
draw_oxide_and_metal(ax2)
ax2.set_xlabel('X (μm)')
ax2.set_ylabel('Y (μm)')
ax2.set_title('2D Acceptor Doping Profile (p11, p12, p13)')
ax2.set_xlim(-W_SIM / um, W_SIM / um)
ax2.set_ylim(H_SI/um + 15.0, -110.0)

plt.tight_layout()
plt.savefig('doping_2d.png', dpi=300)
plt.close()
print("Plot saved to doping_2d.png")