from opentrons.protocol_api import ProtocolContext

# Metadata
metadata = {
    'protocolName': 'Gradient Mixing with Dilution Strategy',
    'author': 'Opentrons AI Assistant',
    'description': 'Creates a 96-well gradient of NaOH and NH4HCO3 with a fixed amount of SmCo Leach. Uses dilutions for small volumes (<10uL).',
    'apiLevel': '2.16'
}
requirements = {"robotType": "Flex"}

def run(protocol: ProtocolContext):
    """
    This protocol performs a gradient mixing experiment in a 96-well plate.
    
    **Procedure Overview:**
    1.  Adds a fixed 200 µL of 'SmCo Leach' to all 96 wells.
    2.  Dispenses variable volumes of DI Water, NaOH, and NH4HCO3 to achieve
        a final volume of 400 µL per well.
    3.  A 1:5 dilution strategy is used for NaOH and NH4HCO3 volumes under 10 µL.
        This script assumes 0.2M versions of these reagents have been pre-made
        and placed in the reservoir.
    """

    # --- 1. Labware & Pipette Setup ---
    protocol.comment("--- Loading Labware and Pipette ---")
    
    # Destination Plate: A deep-well plate is used to safely hold 400uL.
    dest_plate = protocol.load_labware(
        'axygen_96_wellplate_500ul', 'D2', 'Destination Plate'
    )
    trash_bin = protocol.load_trash_bin("A3")
    # Reagent Reservoir
    reservoir = protocol.load_labware(
        '', 'D3', 'Reagent Reservoir'
    )
    
    # Tip Racks for the P1000 pipette
    tiprack_1000_1 = protocol.load_labware('opentrons_flex_96_tiprack_1000ul', 'C1')
    tiprack_1000_2 = protocol.load_labware('opentrons_flex_96_tiprack_1000ul', 'D1')
    
    # Pipette: Using a 1000uL pipette as its min volume is 10uL after dilutions.
    pipette = protocol.load_instrument(
        'flex_1channel_1000', 'left', tip_racks=[tiprack_1000_1, tiprack_1000_2]
    )

    # --- 2. Reagent Location Setup ---
    protocol.comment("--- Defining Reagent Locations ---")
    smco_leach = reservoir.wells_by_name()['A1']
    di_water = reservoir.wells_by_name()['A2']
    naoh_stock = reservoir.wells_by_name()['A3']        # 1M NaOH
    naoh_diluted = reservoir.wells_by_name()['A4']      # 0.2M NaOH (1:5 dilution)
    nh4hco3_stock = reservoir.wells_by_name()['A5']     # 1M NH4HCO3
    nh4hco3_diluted = reservoir.wells_by_name()['A6']   # 0.2M NH4HCO3 (1:5 dilution)

    # --- 3. Volume Definitions ---
    # These are the target volumes of the stock (1M) solutions.
    naoh_vols_by_row = [0, 5, 10, 20, 30, 40, 60, 80]  # For rows A-H
    nh4hco3_vols_by_col = [0, 2, 4, 8, 12, 16, 25, 35, 50, 65, 80, 100]  # For cols 1-12
    DILUTION_FACTOR = 5
    MIN_VOL = 10.0

    # --- 4. Protocol Steps ---

    # Step 1: Dispense SmCo Leach (fixed volume)
    protocol.comment("--- Step 1: Dispensing 200uL SmCo Leach to all wells ---")
    # The distribute command efficiently transfers a fixed volume to multiple wells.
    # It will automatically manage tip usage and refills.
    pipette.distribute(200, smco_leach, dest_plate.wells())

    # Step 2: Dispense DI Water (variable volume)
    protocol.comment("--- Step 2: Dispensing DI Water ---")
    # Water volumes are unique to each well, so we calculate them first.
    water_transfers = []
    for r_idx, row in enumerate(dest_plate.rows()):
        for c_idx, well in enumerate(row):
            # Get target stock volumes
            target_naoh_vol = naoh_vols_by_row[r_idx]
            target_nh4_vol = nh4hco3_vols_by_col[c_idx]

            # Calculate how much extra water will be added via diluted reagents
            water_from_dilutions = 0
            if 0 < target_naoh_vol < MIN_VOL:
                water_from_dilutions += (target_naoh_vol * DILUTION_FACTOR) - target_naoh_vol
            if 0 < target_nh4_vol < MIN_VOL:
                water_from_dilutions += (target_nh4_vol * DILUTION_FACTOR) - target_nh4_vol
            
            # Calculate final DI water volume to add separately
            original_di_water_vol = 100 - target_naoh_vol - target_nh4_vol
            final_di_water_vol = original_di_water_vol - water_from_dilutions

            if final_di_water_vol > 0:
                water_transfers.append({'volume': final_di_water_vol, 'dest': well})

    # Perform the water transfers using one tip for efficiency
    pipette.pick_up_tip()
    for transfer in water_transfers:
        pipette.transfer(transfer['volume'], di_water, transfer['dest'], new_tip='never')
    pipette.drop_tip()

    # Step 3: Dispense NaOH solutions
    protocol.comment("--- Step 3: Dispensing NaOH Solutions ---")
    # Use a new tip for each distribute call to prevent cross-contamination
    
    # 3a: Diluted NaOH for Row B (5 µL target)
    target_vol_b = naoh_vols_by_row[1]
    if 0 < target_vol_b < MIN_VOL:
        vol_to_dispense = target_vol_b * DILUTION_FACTOR  # 5 * 5 = 25 µL
        pipette.distribute(vol_to_dispense, naoh_diluted, dest_plate.rows_by_name()['B'])

    # 3b: Stock NaOH for Rows C-H (>= 10 µL targets)
    stock_naoh_rows = dest_plate.rows()[2:]  # C, D, E, F, G, H
    stock_naoh_vols = naoh_vols_by_row[2:]
    for i, vol in enumerate(stock_naoh_vols):
        if vol >= MIN_VOL:
            pipette.distribute(vol, naoh_stock, stock_naoh_rows[i])

    # Step 4: Dispense NH4HCO3 solutions
    protocol.comment("--- Step 4: Dispensing NH4HCO3 Solutions ---")

    # 4a: Diluted NH4HCO3 for Cols 2, 3, 4 (2, 4, 8 µL targets)
    diluted_nh4_cols = dest_plate.columns()[1:4]  # 2, 3, 4
    diluted_nh4_vols = nh4hco3_vols_by_col[1:4]
    for i, target_vol in enumerate(diluted_nh4_vols):
        if 0 < target_vol < MIN_VOL:
            vol_to_dispense = target_vol * DILUTION_FACTOR
            pipette.distribute(vol_to_dispense, nh4hco3_diluted, diluted_nh4_cols[i])

    # 4b: Stock NH4HCO3 for Cols 5-12 (>= 10 µL targets)
    stock_nh4_cols = dest_plate.columns()[4:]  # 5 through 12
    stock_nh4_vols = nh4hco3_vols_by_col[4:]
    for i, vol in enumerate(stock_nh4_vols):
        if vol >= MIN_VOL:
            pipette.distribute(vol, nh4hco3_stock, stock_nh4_cols[i])
            
    protocol.comment("--- Protocol Complete ---")
