import devsim
import pickle
import dynamic_refine
import sys
import os

OUT_DIR = "output_this_run/"

# 1. Load latest checkpoint/seed
checkpoint_file = f"{OUT_DIR}seed_45V.pkl"
    
if not os.path.exists(checkpoint_file):
    print("Seed file not found. Please make sure the sweep completed.")
    sys.exit(1)

with open(checkpoint_file, "rb") as f:
    cp_data = pickle.load(f)

v_current = cp_data.get("v_current", cp_data.get("voltage"))
state = cp_data["state"]

print(f"Loaded checkpoint at V = {v_current:.2f} V")

# 2. Setup initial device and load mesh
device = "device_2d"
devsim.create_gmsh_mesh(mesh=device, file="device_2d.msh")
dynamic_refine.setup_physics_for_device(device, is_avalanche_enabled=False)

# Restore 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"])

# Apply bias parameters
for c in ["MT1_Si", "MT1_P12_Si", "MT1_Ox", "MT1_Mold"]:
    devsim.set_parameter(device=device, name=f"{c}_bias", value=v_current)
for c in ["MT2_Si", "MT2_P12_Si", "MT2_Ox", "MT2_Mold"]:
    devsim.set_parameter(device=device, name=f"{c}_bias", value=0.0)

# Run Poisson and initial drift-diffusion solve to ensure fields are calculated
devsim.solve(type="dc", absolute_error=1e10, relative_error=1e-3, charge_error=1e12, maximum_iterations=10, info=True)

# 3. Call dynamic refinement and interpolation!
device_refined, opts = dynamic_refine.refine_and_interpolate(device, v_current, is_avalanche_enabled=False)

print("\n--- Test result ---")
print(f"Refinement successful! Refined device name: {device_refined}")
print("New mesh successfully generated and variables interpolated!")