from opentrons import protocol_api

metadata = {
    'protocolName': 'Selective Precipitation Screening',
    'author': 'OpentronsAI',
    'description': 'Screening experiment with multiple reagents and concentrations in 96-well plate',
    'source': 'OpentronsAI'
}

requirements = {
    'robotType': 'Flex',
    'apiLevel': '2.21'
}

def run(protocol: protocol_api.ProtocolContext):
    # Load labware
    plate = protocol.load_labware('nest_96_wellplate_200ul_flat', 'D2')
    reservoir1 = protocol.load_labware('nest_12_reservoir_15ml', 'C2')
    reservoir2 = protocol.load_labware('nest_12_reservoir_15ml', 'D1')
    water_reservoir = protocol.load_labware('nest_1_reservoir_195ml', 'C1')
    
    # Load tip racks
    tiprack1 = protocol.load_labware('opentrons_flex_96_tiprack_1000ul', 'B1')
    tiprack2 = protocol.load_labware('opentrons_flex_96_tiprack_1000ul', 'B2')
    tiprack3 = protocol.load_labware('opentrons_flex_96_tiprack_1000ul', 'A1')
    
    # Load trash
    trash = protocol.load_trash_bin('A3')
    
    # Load pipette
    pipette = protocol.load_instrument(
        'flex_1channel_1000',
        'left',
        tip_racks=[tiprack1, tiprack2, tiprack3]
    )
    
    # Part 1: Dispense 150 µL produced water into all 96 wells
    pipette.pick_up_tip()
    for well in plate.wells():
        pipette.aspirate(150, water_reservoir['A1'])
        pipette.dispense(150, well)
    pipette.drop_tip()
    
    # Part 2: Add reagents according to plate layout
    
    # NaOH additions from reservoir1
    # Wells A1,A2,A3: 15 µL from A1 (NaOH 0.10M)
    pipette.transfer(15, reservoir1['A1'], [plate['A1'], plate['A2'], plate['A3']], new_tip='once')
    
    # Wells A4,A5,A6: 50 µL from A1 (NaOH 0.10M)
    pipette.transfer(50, reservoir1['A1'], [plate['A4'], plate['A5'], plate['A6']], new_tip='once')
    
    # Wells A7,A8,A9: 15 µL from A2 (NaOH 0.50M)
    pipette.transfer(15, reservoir1['A2'], [plate['A7'], plate['A8'], plate['A9']], new_tip='once')
    
    # Wells A10,A11,A12: 50 µL from A2 (NaOH 0.50M)
    pipette.transfer(50, reservoir1['A2'], [plate['A10'], plate['A11'], plate['A12']], new_tip='once')
    
    # Wells B1,B2,B3: 15 µL from A3 (NaOH 1.00M)
    pipette.transfer(15, reservoir1['A3'], [plate['B1'], plate['B2'], plate['B3']], new_tip='once')
    
    # Wells B4,B5,B6: 50 µL from A3 (NaOH 1.00M)
    pipette.transfer(50, reservoir1['A3'], [plate['B4'], plate['B5'], plate['B6']], new_tip='once')
    
    # Na2CO3 additions from reservoir1
    # Wells B7,B8,B9: 15 µL from A4 (Na2CO3 0.10M)
    pipette.transfer(15, reservoir1['A4'], [plate['B7'], plate['B8'], plate['B9']], new_tip='once')
    
    # Wells B10,B11,B12: 50 µL from A4 (Na2CO3 0.10M)
    pipette.transfer(50, reservoir1['A4'], [plate['B10'], plate['B11'], plate['B12']], new_tip='once')
    
    # Wells C1,C2,C3: 15 µL from A5 (Na2CO3 0.50M)
    pipette.transfer(15, reservoir1['A5'], [plate['C1'], plate['C2'], plate['C3']], new_tip='once')
    
    # Wells C4,C5,C6: 50 µL from A5 (Na2CO3 0.50M)
    pipette.transfer(50, reservoir1['A5'], [plate['C4'], plate['C5'], plate['C6']], new_tip='once')
    
    # Wells C7,C8,C9: 15 µL from A6 (Na2CO3 1.00M)
    pipette.transfer(15, reservoir1['A6'], [plate['C7'], plate['C8'], plate['C9']], new_tip='once')
    
    # Wells C10,C11,C12: 50 µL from A6 (Na2CO3 1.00M)
    pipette.transfer(50, reservoir1['A6'], [plate['C10'], plate['C11'], plate['C12']], new_tip='once')
    
    # Na3PO4 additions from reservoir1
    # Wells D1,D2,D3: 15 µL from A7 (Na3PO4 0.05M)
    pipette.transfer(15, reservoir1['A7'], [plate['D1'], plate['D2'], plate['D3']], new_tip='once')
    
    # Wells D4,D5,D6: 50 µL from A7 (Na3PO4 0.05M)
    pipette.transfer(50, reservoir1['A7'], [plate['D4'], plate['D5'], plate['D6']], new_tip='once')
    
    # Wells D7,D8,D9: 15 µL from A8 (Na3PO4 0.25M)
    pipette.transfer(15, reservoir1['A8'], [plate['D7'], plate['D8'], plate['D9']], new_tip='once')
    
    # Wells D10,D11,D12: 50 µL from A8 (Na3PO4 0.25M)
    pipette.transfer(50, reservoir1['A8'], [plate['D10'], plate['D11'], plate['D12']], new_tip='once')
    
    # Wells E1,E2,E3: 15 µL from A9 (Na3PO4 0.50M)
    pipette.transfer(15, reservoir1['A9'], [plate['E1'], plate['E2'], plate['E3']], new_tip='once')
    
    # Wells E4,E5,E6: 50 µL from A9 (Na3PO4 0.50M)
    pipette.transfer(50, reservoir1['A9'], [plate['E4'], plate['E5'], plate['E6']], new_tip='once')
    
    # NaF additions from reservoir1
    # Wells E7,E8,E9: 15 µL from A10 (NaF 0.05M)
    pipette.transfer(15, reservoir1['A10'], [plate['E7'], plate['E8'], plate['E9']], new_tip='once')
    
    # Wells E10,E11,E12: 50 µL from A10 (NaF 0.05M)
    pipette.transfer(50, reservoir1['A10'], [plate['E10'], plate['E11'], plate['E12']], new_tip='once')
    
    # Wells F1,F2,F3: 15 µL from A11 (NaF 0.25M)
    pipette.transfer(15, reservoir1['A11'], [plate['F1'], plate['F2'], plate['F3']], new_tip='once')
    
    # Wells F4,F5,F6: 50 µL from A11 (NaF 0.25M)
    pipette.transfer(50, reservoir1['A11'], [plate['F4'], plate['F5'], plate['F6']], new_tip='once')
    
    # Wells F7,F8,F9: 15 µL from A12 (NaF 0.50M)
    pipette.transfer(15, reservoir1['A12'], [plate['F7'], plate['F8'], plate['F9']], new_tip='once')
    
    # Wells F10,F11,F12: 50 µL from A12 (NaF 0.50M)
    pipette.transfer(50, reservoir1['A12'], [plate['F10'], plate['F11'], plate['F12']], new_tip='once')
    
    # Oxalic acid additions from reservoir2
    # Wells G1,G2,G3: 15 µL from A1 (Oxalic 0.05M)
    pipette.transfer(15, reservoir2['A1'], [plate['G1'], plate['G2'], plate['G3']], new_tip='once')
    
    # Wells G4,G5,G6: 50 µL from A1 (Oxalic 0.05M)
    pipette.transfer(50, reservoir2['A1'], [plate['G4'], plate['G5'], plate['G6']], new_tip='once')
    
    # Wells G7,G8,G9: 15 µL from A2 (Oxalic 0.25M)
    pipette.transfer(15, reservoir2['A2'], [plate['G7'], plate['G8'], plate['G9']], new_tip='once')
    
    # Wells G10,G11,G12: 50 µL from A2 (Oxalic 0.25M)
    pipette.transfer(50, reservoir2['A2'], [plate['G10'], plate['G11'], plate['G12']], new_tip='once')
    
    # Wells H1,H2,H3: 15 µL from A3 (Oxalic 0.50M)
    pipette.transfer(15, reservoir2['A3'], [plate['H1'], plate['H2'], plate['H3']], new_tip='once')
    
    # Wells H4,H5,H6: 50 µL from A3 (Oxalic 0.50M)
    pipette.transfer(50, reservoir2['A3'], [plate['H4'], plate['H5'], plate['H6']], new_tip='once')
    
    # Control additions (DI water from reservoir2 well A4)
    # Wells H7,H8,H9,H10,H11,H12: 50 µL from A4
    pipette.transfer(50, reservoir2['A4'], [plate['H7'], plate['H8'], plate['H9'], 
                                             plate['H10'], plate['H11'], plate['H12']], new_tip='once')
    
    # Part 3: Mix all wells
    # Group wells by reagent type to minimize cross-contamination
    
    # Mix NaOH wells (A1-B6)
    naoh_wells = [plate[well] for well in ['A1', 'A2', 'A3', 'A4', 'A5', 'A6',
                                             'A7', 'A8', 'A9', 'A10', 'A11', 'A12',
                                             'B1', 'B2', 'B3', 'B4', 'B5', 'B6']]
    pipette.pick_up_tip()
    for well in naoh_wells:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Mix Na2CO3 wells (B7-C12)
    na2co3_wells = [plate[well] for well in ['B7', 'B8', 'B9', 'B10', 'B11', 'B12',
                                               'C1', 'C2', 'C3', 'C4', 'C5', 'C6',
                                               'C7', 'C8', 'C9', 'C10', 'C11', 'C12']]
    pipette.pick_up_tip()
    for well in na2co3_wells:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Mix Na3PO4 wells (D1-E6)
    na3po4_wells = [plate[well] for well in ['D1', 'D2', 'D3', 'D4', 'D5', 'D6',
                                               'D7', 'D8', 'D9', 'D10', 'D11', 'D12',
                                               'E1', 'E2', 'E3', 'E4', 'E5', 'E6']]
    pipette.pick_up_tip()
    for well in na3po4_wells:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Mix NaF wells (E7-F12)
    naf_wells = [plate[well] for well in ['E7', 'E8', 'E9', 'E10', 'E11', 'E12',
                                            'F1', 'F2', 'F3', 'F4', 'F5', 'F6',
                                            'F7', 'F8', 'F9', 'F10', 'F11', 'F12']]
    pipette.pick_up_tip()
    for well in naf_wells:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Mix Oxalic acid wells (G1-H6)
    oxalic_wells = [plate[well] for well in ['G1', 'G2', 'G3', 'G4', 'G5', 'G6',
                                               'G7', 'G8', 'G9', 'G10', 'G11', 'G12',
                                               'H1', 'H2', 'H3', 'H4', 'H5', 'H6']]
    pipette.pick_up_tip()
    for well in oxalic_wells:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Mix control wells (H7-H12)
    control_wells = [plate[well] for well in ['H7', 'H8', 'H9', 'H10', 'H11', 'H12']]
    pipette.pick_up_tip()
    for well in control_wells:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
