from opentrons import protocol_api

metadata = {
    'protocolName': 'Round 2 Sm Precipitation Optimization',
    'author': 'Research Team',
    'description': 'NaOH-Bicarbonate gradient for Sm recovery optimization',
    'apiLevel': '2.15'
}

def run(protocol: protocol_api.ProtocolContext):
    # Load trash bin at position A3 (required for Flex)
    trash = protocol.load_trash_bin('A3')
    
    # Load labware
    reservoir = protocol.load_labware('nest_12_reservoir_15ml', 'C1')
    plate = protocol.load_labware('axygen_96_wellplate_500ul', 'D1')
    tiprack_1000 = protocol.load_labware('opentrons_flex_96_tiprack_1000ul', 'C2')
    tiprack_200 = protocol.load_labware('opentrons_flex_96_tiprack_200ul', 'C3')
    
    # Load pipette - Flex 1-channel 1000 µL on left mount
    pipette = protocol.load_instrument(
        'flex_1channel_1000', 
        'left', 
        tip_racks=[tiprack_1000, tiprack_200]
    )
    
    # Define reagent locations in reservoir
    naoh = reservoir['A1']          # 1.0M NaOH solution
    bicarbonate = reservoir['A2']   # 1.0M Bicarbonate solution
    water = reservoir['A3']         # Deionized water
    leachate = reservoir['A4']      # SmCo leachate starting material
    
    # Define NaOH volumes by row (A=0, B=1, ..., H=7)
    # Row A: 0 µL, Row B: 1.4 µL, Row C: 2.9 µL, Row D: 4.3 µL
    # Row E: 5.7 µL, Row F: 7.1 µL, Row G: 8.6 µL, Row H: 10 µL
    naoh_volumes = [0, 1.4, 2.9, 4.3, 5.7, 7.1, 8.6, 10.0]
    
    # Define Bicarbonate volumes by column (1-12)
    # Column 1: 0 µL, Column 2: 3 µL, Column 3: 4 µL, Column 4: 6 µL
    # Column 5: 8 µL, Column 6: 10 µL, Column 7: 12 µL, Column 8: 15 µL
    # Column 9: 18 µL, Column 10: 20 µL, Column 11: 25 µL, Column 12: 30 µL
    bicarbonate_volumes = [0, 3, 4, 6, 8, 10, 12, 15, 18, 20, 25, 30]
    
    # Rows and columns for iteration
    rows = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']
    columns = range(1, 13)  # 1 to 12
    
    # Constants
    LEACHATE_VOLUME = 200  # µL of SmCo leachate per well
    TOTAL_VOLUME = 400     # µL total volume per well
    BASE_REAGENT_VOLUME = 200  # µL allocated for reagents (NaOH + Bicarbonate + Water)
    
    # Pre-calculate water volumes for all wells
    water_volumes = {}
    for row_idx, row in enumerate(rows):
        for col_idx, col in enumerate(columns):
            well_name = f'{row}{col}'
            naoh_vol = naoh_volumes[row_idx]
            bicarb_vol = bicarbonate_volumes[col_idx]
            
            # Calculate water volume to bring total to 400 µL
            # Water volume = 200 (base reagent allocation) - NaOH - Bicarbonate
            water_vol = BASE_REAGENT_VOLUME - naoh_vol - bicarb_vol
            
            # Ensure water volume is not negative
            if water_vol < 0:
                water_vol = 0
                protocol.comment(f"Warning: Well {well_name} - water volume calculated as negative, set to 0")
            
            water_volumes[well_name] = water_vol
    
    # Step 1: Dispense leachate (200 µL) to all wells using 1000 µL tips
    protocol.comment("Step 1: Dispensing leachate to all wells")
    for row in rows:
        for col in columns:
            well = plate[f'{row}{col}']
            pipette.pick_up_tip(tiprack_1000)
            pipette.aspirate(LEACHATE_VOLUME, leachate)
            pipette.dispense(LEACHATE_VOLUME, well)
            pipette.drop_tip()
    
    # Step 2: Dispense water to all wells (where volume > 0) using 1000 µL tips
    protocol.comment("Step 2: Dispensing water to all wells")
    for row in rows:
        for col in columns:
            well_name = f'{row}{col}'
            well = plate[well_name]
            water_vol = water_volumes[well_name]
            
            if water_vol > 0:
                pipette.pick_up_tip(tiprack_1000)
                pipette.aspirate(water_vol, water)
                pipette.dispense(water_vol, well)
                pipette.drop_tip()
    
    # Step 3: Dispense NaOH to all wells (where volume > 0) using 200 µL tips for accuracy
    protocol.comment("Step 3: Dispensing NaOH to all wells (using 200 µL tips for small volumes)")
    for row_idx, row in enumerate(rows):
        naoh_vol = naoh_volumes[row_idx]
        
        if naoh_vol > 0:
            for col in columns:
                well = plate[f'{row}{col}']
                pipette.pick_up_tip(tiprack_200)
                pipette.aspirate(naoh_vol, naoh)
                pipette.dispense(naoh_vol, well)
                pipette.drop_tip()
    
    # Step 4: Dispense Bicarbonate to all wells (where volume > 0) using 200 µL tips for accuracy
    protocol.comment("Step 4: Dispensing bicarbonate to all wells (using 200 µL tips for small volumes)")
    for col_idx, col in enumerate(columns):
        bicarb_vol = bicarbonate_volumes[col_idx]
        
        if bicarb_vol > 0:
            for row in rows:
                well = plate[f'{row}{col}']
                pipette.pick_up_tip(tiprack_200)
                pipette.aspirate(bicarb_vol, bicarbonate)
                pipette.dispense(bicarb_vol, well)
                pipette.drop_tip()
    
    # Protocol complete - provide summary comment
    protocol.comment("Protocol complete. Allow precipitation for 30 minutes at room temperature.")
    protocol.comment("Then centrifuge plate or allow precipitate to settle.")
    protocol.comment("Collect supernatant and precipitate fractions for ICP-MS analysis.")