import json
from opentrons.protocol_api import ProtocolContext

# metadata
metadata = {
    'protocolName': 'Nd Purification via Selective Oxalate Precipitation',
    'author': 'Research Software Engineer',
    'description': 'A protocol to screen pH and precipitant concentrations for purifying an enriched Nd solution.',
    'apiLevel': '2.16'
}

# requirements
requirements = {
    "robotType": "Flex",
}

def run(protocol: ProtocolContext):
    """
    This protocol executes a purification experiment for an enriched Neodymium solution.
    1. Dispenses a fixed volume of 'Enriched Nd Stock' into a 96-well deep-well plate.
    2. Creates a pH gradient across the rows by adding variable volumes of HCl and NaOH.
    3. Creates a precipitant gradient across columns by adding variable volumes of Sodium Oxalate.
    4. Pauses for incubation and subsequent manual steps.
    """
    
    # --- Labware & Module Setup ---
    protocol.comment("--- Loading Labware, Modules, and Pipettes ---")
    
    # Labware
    deepwell_plate = protocol.load_labware(
        "nest_96_wellplate_2ml_deep", "D1", "Deep Well Plate"
    )
    
    reservoir = protocol.load_labware('nest_12_reservoir_15ml', 'C2', "Reagent Reservoir")
    
    tiprack_1000 = protocol.load_labware('opentrons_flex_96_tiprack_1000ul', 'D2')

    # Pipette
    pipette_1000 = protocol.load_instrument('flex_1channel_1000', 'left', tip_racks=[tiprack_1000])

    # Trash
    trash_bin = protocol.load_trash_bin("A3")

    # --- Reagent Locations ---
    protocol.comment("--- Defining Reagent Locations ---")
    enriched_nd_stock = reservoir.wells_by_name()['A1']
    hcl_stock = reservoir.wells_by_name()['A2']
    naoh_stock = reservoir.wells_by_name()['A3']
    na_oxalate_stock = reservoir.wells_by_name()['A4']
    
    # Add instructions for the operator in the run log
    protocol.comment("Setup: Place the following reagents in the C2 Reservoir:")
    protocol.comment("'Enriched Nd Stock' (in 1M HCl) -> Well A1")
    protocol.comment("0.5M HCl -> Well A2")
    protocol.comment("0.5M NaOH -> Well A3")
    protocol.comment("0.25M Sodium Oxalate -> Well A4")

    # --- Step 1: Dispense 'Enriched Nd Stock' (Experimental Step 5) ---
    protocol.comment("--- Dispensing 100 µL of Enriched Nd Stock to all wells ---")
    pipette_1000.pick_up_tip()
    pipette_1000.distribute(
        100,
        enriched_nd_stock,
        deepwell_plate.wells(),
        new_tip='never'
    )
    pipette_1000.drop_tip()

    # --- Step 2: Create pH Gradient (Experimental Step 6) ---
    protocol.comment("--- Creating pH gradient across rows ---")
    
    hcl_volumes_by_row = {'A': 80, 'B': 40, 'C': 20}
    protocol.comment("Dispensing 0.5M HCl to rows A-C.")
    pipette_1000.pick_up_tip()
    for row_name, volume in hcl_volumes_by_row.items():
        if volume > 0:
            pipette_1000.distribute(volume, hcl_stock, deepwell_plate.rows_by_name()[row_name], new_tip='never')
    pipette_1000.drop_tip()

    naoh_volumes_by_row = {'F': 20, 'G': 40, 'H': 80}
    protocol.comment("Dispensing 0.5M NaOH to rows F-H.")
    pipette_1000.pick_up_tip()
    for row_name, volume in naoh_volumes_by_row.items():
        if volume > 0:
            pipette_1000.distribute(volume, naoh_stock, deepwell_plate.rows_by_name()[row_name], new_tip='never')
    pipette_1000.drop_tip()

    # --- Step 3: Create Precipitant Gradient (Experimental Step 7) ---
    protocol.comment("--- Creating precipitant gradient across columns ---")
    
    oxalate_volumes_by_col = {str(i + 1): (i * 20) + 20 for i in range(12)}
    protocol.comment("Dispensing 0.25M Sodium Oxalate to columns 1-12.")
    pipette_1000.pick_up_tip()
    for col_name, volume in oxalate_volumes_by_col.items():
        pipette_1000.distribute(volume, na_oxalate_stock, deepwell_plate.columns_by_name()[col_name], new_tip='never')
    pipette_1000.drop_tip()

    # --- Step 4: Incubation (Experimental Step 8) ---
    protocol.comment("--- Incubating at room temperature for 2 hours ---")
    protocol.delay(hours=2)

    # --- Step 5: Pause for Manual Intervention (Experimental Steps 9 & 10) ---
    protocol.pause(
        "Incubation complete. Please remove the plate from the deck. "
        "Centrifuge the plate to pellet the precipitate, then decant the supernatant. "
        "The precipitates are now ready for analysis (e.g., ICP-MS or EDS)."
    )

    protocol.comment("Protocol complete.")