import json
from opentrons.protocol_api import ProtocolContext

# metadata
metadata = {
    'protocolName': 'Rare Earth Element Precipitation Screening',
    'author': 'Research Software Engineer',
    'description': 'Automated precipitation of Nd, Pr, and Fe from leachate using NaOH and Sodium Oxalate.',
}

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

def run(protocol: ProtocolContext):
    """
    This protocol performs a precipitation screening experiment for rare earth elements.
    1. Dispenses NdFeB leachate into a 96-well plate with a row-wise volume gradient.
    2. Dispenses NaOH and Sodium Oxalate with a column-wise volume gradient.
    3. Shakes the plate to promote precipitation.
    4. Pauses for manual centrifugation and supernatant collection.
    """
    
    # --- Labware & Module Setup ---
    protocol.comment("--- Loading Labware, Modules, and Pipettes ---")
    
    # Modules
    plate = protocol.load_labware(
        "axygen_96_wellplate_500ul",
        location="D1",
        namespace="opentrons",
        version=2,
    )

    # Reservoir for reagents
    reservoir = protocol.load_labware('nest_12_reservoir_15ml', 'C2')
    
    # Tip racks
    tiprack = protocol.load_labware('opentrons_flex_96_tiprack_1000ul', 'D2')

    # Trash
    trash_bin_1 = protocol.load_trash_bin("A3")
    
    # Pipette
    pipette = protocol.load_instrument('flex_1channel_1000', 'left', tip_racks=[tiprack])

    # --- Reagent Locations ---
    protocol.comment("--- Defining Reagent Locations ---")
    leachate_solution = reservoir.wells_by_name()['A1']
    naoh_solution = reservoir.wells_by_name()['A2']
    oxalate_solution = reservoir.wells_by_name()['A3']
    
    # Add instructions for the operator in the run log
    protocol.comment("Place 'D3 100k NdFeB leachate' in Reservoir Well A1.")
    protocol.comment("Place '1M NaOH' in Reservoir Well A2.")
    protocol.comment("Place '0.2M Sodium Oxalate' in Reservoir Well A3.")

    # --- Step 2: Dispense NdFeB Leachate ---
    protocol.comment("--- Dispensing NdFeB Leachate ---")
    leachate_volumes_by_row = {
        'A': 40, 'B': 50, 'C': 60, 'D': 70,
        'E': 80, 'F': 90, 'G': 100, 'H': 110
    }
    
    for row_name, volume in leachate_volumes_by_row.items():
        protocol.comment(f"Dispensing {volume} uL of leachate to Row {row_name}")
        pipette.distribute(
            volume,
            leachate_solution,
            plate.rows_by_name()[row_name],
            new_tip='once'  # Use one tip per row for efficiency
        )
    
    # --- Step 3: Dispense NaOH ---
    protocol.comment("--- Dispensing NaOH Solution ---")
    naoh_volumes_by_col = {
        '1': 30, '2': 60, '3': 90,
        '4': 120, '5': 150, '6': 180
    }

    for col_name, volume in naoh_volumes_by_col.items():
        protocol.comment(f"Dispensing {volume} uL of NaOH to Column {col_name}")
        pipette.distribute(
            volume,
            naoh_solution,
            plate.columns_by_name()[col_name],
            new_tip='once'  # Use one tip per column
        )

    # --- Step 4: Dispense Sodium Oxalate ---
    protocol.comment("--- Dispensing Sodium Oxalate Solution ---")
    oxalate_volumes_by_col = {
        '7': 30, '8': 60, '9': 90,
        '10': 120, '11': 150, '12': 180
    }

    for col_name, volume in oxalate_volumes_by_col.items():
        protocol.comment(f"Dispensing {volume} uL of Sodium Oxalate to Column {col_name}")
        pipette.distribute(
            volume,
            oxalate_solution,
            plate.columns_by_name()[col_name],
            new_tip='once'  # Use one tip per column
        )
    


    
    protocol.comment("Protocol complete. The robot is ready for the next run.")