import re

with open('solve_sweep_2d.py', 'r') as f:
    code = f.read()

# Make solve_sweep_recon.py
recon_code = code

# 1. Add import pickle
recon_code = "import pickle\n" + recon_code

# 2. Modify equations setup to NOT have Avalanche by default
eq_setup = """# 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)"""

recon_code = re.sub(r'# Instantiate Avalanche.*?min_error=1e5\)', eq_setup, recon_code, flags=re.DOTALL)

# 3. Add Recon logic inside the sweep loop
recon_logic = """
        # --- RECONNAISSANCE PROBE ---
        if v_current >= next_recon_v:
            state_data = save_state(device)
            seed_data = {"voltage": v_current, "step_size": step_size, "state": state_data}
            seed_filename = f"seed_{int(next_recon_v)}V.pkl"
            with open(seed_filename, "wb") as f:
                pickle.dump(seed_data, f)
            print(f"\\n--- RECON PROBE at {v_current:.2f} V ---")
            print(f"Saved seed to {seed_filename}")
            
            # Turn ON Avalanche
            devsim.equation(device=device, region="Silicon", name="ElectronContinuityEquation", variable_name="Electrons",
                            time_node_model="NCharge", edge_model=opts['Jn'], edge_volume_model="AvalancheGeneration", 
                            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="AvalancheGeneration", 
                            variable_update="positive", node_model="HoleGeneration", min_error=1e5)
            
            try:
                # Stage 1 pre-conditioning
                try:
                    devsim.solve(type="dc", absolute_error=1e10, relative_error=1e-1, charge_error=1e12, maximum_iterations=10, info=True)
                except devsim.error:
                    pass
                # Stage 2 solve
                res_av = devsim.solve(type="dc", absolute_error=1e10, relative_error=1e-3, charge_error=1e12, maximum_iterations=15, info=True)
                if res_av.get("converged", False):
                    # Measure Avalanche current
                    ia_n_si = devsim.get_contact_current(device=device, contact="MT1_Si", equation="ElectronContinuityEquation")
                    ia_p_si = devsim.get_contact_current(device=device, contact="MT1_Si", equation="HoleContinuityEquation")
                    ia_n_p12 = devsim.get_contact_current(device=device, contact="MT1_P12_Si", equation="ElectronContinuityEquation")
                    ia_p_p12 = devsim.get_contact_current(device=device, contact="MT1_P12_Si", equation="HoleContinuityEquation")
                    av_curr = ia_n_si + ia_p_si + ia_n_p12 + ia_p_p12
                    print(f"Avalanche Current at {v_current:.2f} V: {av_curr:.4e} A")
                    with open("recon_avalanche.log", "a") as f:
                        f.write(f"{v_current:.2f}\\t{av_curr:.4e}\\n")
                else:
                    print("Avalanche failed to converge.")
                    with open("recon_avalanche.log", "a") as f:
                        f.write(f"{v_current:.2f}\\tFAILED\\n")
            except devsim.error:
                print("Avalanche failed to converge.")
                with open("recon_avalanche.log", "a") as f:
                    f.write(f"{v_current:.2f}\\tFAILED\\n")
            
            # Restore state and Turn OFF Avalanche
            restore_state(device, state_data)
            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)
            print("--- END RECON PROBE ---\\n")
            next_recon_v += 50.0
            
        # Grow step size for next step adaptively based on Newton iterations
"""

# Insert variables for recon
init_vars = """# Arrays to store I-V data
voltage_list = [0.0]
current_list = [0.0]

# Recon variables
next_recon_v = 50.0
with open("recon_avalanche.log", "w") as f:
    f.write("Voltage(V)\\tAvalancheCurrent(A)\\n")
"""

recon_code = recon_code.replace("# Arrays to store I-V data\nvoltage_list = [0.0]\ncurrent_list = [0.0]", init_vars)
recon_code = recon_code.replace("        # Grow step size for next step adaptively based on Newton iterations", recon_logic)

with open('solve_sweep_recon.py', 'w') as f:
    f.write(recon_code)


# Make solve_sweep_bv.py
bv_code = code
bv_code = "import pickle\nimport argparse\n" + bv_code

# Turn ON Avalanche by default in solve_sweep_bv.py
bv_code = bv_code.replace('edge_volume_model=""', 'edge_volume_model="AvalancheGeneration"')

# Keep Avalanche ON by default
arg_parse = """
parser = argparse.ArgumentParser()
parser.add_argument("--seed", type=str, required=True, help="Seed pickle file to load")
args = parser.parse_args()

print(f"Loading seed state from {args.seed}...")
with open(args.seed, "rb") as f:
    seed_data = pickle.load(f)

# Arrays to store I-V data
voltage_list = [seed_data["voltage"]]
current_list = [0.0] # Will be updated after first solve

# Overwrite initial state variables
v_current = seed_data["voltage"]
step_size = min(0.5, seed_data["step_size"]) # Start small for BV

# Apply biases BEFORE solving
for c in ["MT1_Si", "MT1_P12_Si", "MT1_Ox", "MT1_Mold"]:
    devsim.set_parameter(device=device, name=f"{c}_bias", value=v_current)

restore_state(device, seed_data["state"])
state = save_state(device)
"""

bv_code = bv_code.replace("# Arrays to store I-V data\nvoltage_list = [0.0]\ncurrent_list = [0.0]", arg_parse)

# Remove the initial Poisson and initial DD solve because we are loading from a seed!
# But wait! DEVSIM requires us to run devsim.solve AT LEAST ONCE to initialize the equations?
# Actually, setting the state is enough. But the initial solves are between line 173 and 220.
# We can just let it run the zero bias solve, and then overwrite the state. It takes a few seconds.
# So we don't need to delete the initial solve. It's safer to let it build the matrices.

with open('solve_sweep_bv.py', 'w') as f:
    f.write(bv_code)