from opentrons import protocol_api

metadata = {
    'protocolName': 'Mg-Focused Selective Precipitation Screening',
    'author': 'OpentronsAI',
    'description': 'Two-phase precipitation screening with NaOH titration and sequential oxalic acid + NaOH additions',
    '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')  # NaOH stocks
    reservoir2 = protocol.load_labware('nest_12_reservoir_15ml', 'D1')  # Oxalic acid stocks
    produced_water = 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 produced water to all wells
    # =========================================================================
    protocol.comment("=== PART 1: Dispensing 150 uL produced water to all 96 wells ===")
    pipette.pick_up_tip()
    for well in plate.wells():
        pipette.transfer(150, produced_water['A1'], well, new_tip='never')
    pipette.drop_tip()
    
    # =========================================================================
    # PART 2: Phase A - NaOH reagent additions (Rows A-E)
    # =========================================================================
    protocol.comment("=== PART 2: Phase A - NaOH additions ===")
    
    # Reservoir1 layout:
    # A1: NaOH 0.10M, A2: NaOH 0.25M, A3: NaOH 0.50M, A4: NaOH 0.75M
    # A5: NaOH 1.00M, A6: NaOH 1.50M, A7: NaOH 2.00M, A8: DI water
    
    # Row A: NaOH 0.10M and 0.25M
    pipette.transfer(20, reservoir1['A1'], [plate[w] for w in ['A1', 'A2', 'A3']], new_tip='once')
    pipette.transfer(50, reservoir1['A1'], [plate[w] for w in ['A4', 'A5', 'A6']], new_tip='once')
    pipette.transfer(20, reservoir1['A2'], [plate[w] for w in ['A7', 'A8', 'A9']], new_tip='once')
    pipette.transfer(50, reservoir1['A2'], [plate[w] for w in ['A10', 'A11', 'A12']], new_tip='once')
    
    # Row B: NaOH 0.50M and 0.75M
    pipette.transfer(20, reservoir1['A3'], [plate[w] for w in ['B1', 'B2', 'B3']], new_tip='once')
    pipette.transfer(35, reservoir1['A3'], [plate[w] for w in ['B4', 'B5', 'B6']], new_tip='once')
    pipette.transfer(50, reservoir1['A3'], [plate[w] for w in ['B7', 'B8', 'B9']], new_tip='once')
    pipette.transfer(20, reservoir1['A4'], [plate[w] for w in ['B10', 'B11', 'B12']], new_tip='once')
    
    # Row C: NaOH 0.75M and 1.00M
    pipette.transfer(35, reservoir1['A4'], [plate[w] for w in ['C1', 'C2', 'C3']], new_tip='once')
    pipette.transfer(50, reservoir1['A4'], [plate[w] for w in ['C4', 'C5', 'C6']], new_tip='once')
    pipette.transfer(20, reservoir1['A5'], [plate[w] for w in ['C7', 'C8', 'C9']], new_tip='once')
    pipette.transfer(35, reservoir1['A5'], [plate[w] for w in ['C10', 'C11', 'C12']], new_tip='once')
    
    # Row D: NaOH 1.00M, 1.50M, 2.00M
    pipette.transfer(50, reservoir1['A5'], [plate[w] for w in ['D1', 'D2', 'D3']], new_tip='once')
    pipette.transfer(20, reservoir1['A6'], [plate[w] for w in ['D4', 'D5', 'D6']], new_tip='once')
    pipette.transfer(50, reservoir1['A6'], [plate[w] for w in ['D7', 'D8', 'D9']], new_tip='once')
    pipette.transfer(20, reservoir1['A7'], [plate[w] for w in ['D10', 'D11', 'D12']], new_tip='once')
    
    # Row E: NaOH 2.00M and controls
    pipette.transfer(50, reservoir1['A7'], [plate[w] for w in ['E1', 'E2', 'E3']], new_tip='once')
    pipette.transfer(50, reservoir1['A8'], [plate[w] for w in ['E4', 'E5', 'E6']], new_tip='once')
    pipette.transfer(50, reservoir1['A8'], [plate[w] for w in ['E7', 'E8', 'E9']], new_tip='once')
    
    # =========================================================================
    # PART 3: Phase A mixing (grouped by NaOH concentration)
    # =========================================================================
    protocol.comment("=== PART 3: Phase A mixing ===")
    
    # Group 1: NaOH 0.10M (A1-A6)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['A1', 'A2', 'A3', 'A4', 'A5', 'A6']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 2: NaOH 0.25M (A7-A12)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['A7', 'A8', 'A9', 'A10', 'A11', 'A12']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 3: NaOH 0.50M (B1-B9)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['B1', 'B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B9']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 4: NaOH 0.75M (B10-B12, C1-C6)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['B10', 'B11', 'B12', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 5: NaOH 1.00M (C7-C12, D1-D3)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['C7', 'C8', 'C9', 'C10', 'C11', 'C12', 'D1', 'D2', 'D3']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 6: NaOH 1.50M (D4-D9)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['D4', 'D5', 'D6', 'D7', 'D8', 'D9']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 7: NaOH 2.00M (D10-D12, E1-E3)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['D10', 'D11', 'D12', 'E1', 'E2', 'E3']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 8: Controls (E4-E9)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['E4', 'E5', 'E6', 'E7', 'E8', 'E9']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # =========================================================================
    # PART 4: Phase B Step 1 - Oxalic acid additions (Rows F-H)
    # =========================================================================
    protocol.comment("=== PART 4: Phase B Step 1 - Oxalic acid additions ===")
    
    # Reservoir2 layout:
    # A1: Oxalic acid 0.10M, A2: Oxalic acid 0.25M, A3: Oxalic acid 0.50M, A4: DI water
    
    # Oxalic acid 0.10M (F1-F9)
    pipette.transfer(30, reservoir2['A1'], [plate[w] for w in ['F1', 'F2', 'F3']], new_tip='once')
    pipette.transfer(30, reservoir2['A1'], [plate[w] for w in ['F4', 'F5', 'F6']], new_tip='once')
    pipette.transfer(30, reservoir2['A1'], [plate[w] for w in ['F7', 'F8', 'F9']], new_tip='once')
    
    # Oxalic acid 0.25M (F10-F12, G1-G6, H4-H6)
    pipette.transfer(30, reservoir2['A2'], [plate[w] for w in ['F10', 'F11', 'F12']], new_tip='once')
    pipette.transfer(30, reservoir2['A2'], [plate[w] for w in ['G1', 'G2', 'G3']], new_tip='once')
    pipette.transfer(30, reservoir2['A2'], [plate[w] for w in ['G4', 'G5', 'G6']], new_tip='once')
    pipette.transfer(30, reservoir2['A2'], [plate[w] for w in ['H4', 'H5', 'H6']], new_tip='once')
    
    # Oxalic acid 0.50M (G7-G12, H1-H3)
    pipette.transfer(30, reservoir2['A3'], [plate[w] for w in ['G7', 'G8', 'G9']], new_tip='once')
    pipette.transfer(30, reservoir2['A3'], [plate[w] for w in ['G10', 'G11', 'G12']], new_tip='once')
    pipette.transfer(30, reservoir2['A3'], [plate[w] for w in ['H1', 'H2', 'H3']], new_tip='once')
    
    # DI water controls (H7-H12)
    pipette.transfer(30, reservoir2['A4'], [plate[w] for w in ['H7', 'H8', 'H9']], new_tip='once')
    pipette.transfer(30, reservoir2['A4'], [plate[w] for w in ['H10', 'H11', 'H12']], new_tip='once')
    
    # =========================================================================
    # PART 5: Phase B Step 1 mixing (after oxalic acid)
    # =========================================================================
    protocol.comment("=== PART 5: Phase B Step 1 mixing ===")
    
    # Group 1: Oxalic 0.10M (F1-F9)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['F1', 'F2', 'F3', 'F4', 'F5', 'F6', 'F7', 'F8', 'F9']]:
        pipette.mix(3, 100, well)
    pipette.drop_tip()
    
    # Group 2: Oxalic 0.25M (F10-F12, G1-G6, H4-H6)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['F10', 'F11', 'F12', 'G1', 'G2', 'G3', 'G4', 'G5', 'G6', 'H4', 'H5', 'H6']]:
        pipette.mix(3, 100, well)
    pipette.drop_tip()
    
    # Group 3: Oxalic 0.50M (G7-G12, H1-H3)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['G7', 'G8', 'G9', 'G10', 'G11', 'G12', 'H1', 'H2', 'H3']]:
        pipette.mix(3, 100, well)
    pipette.drop_tip()
    
    # Group 4: Controls (H7-H12)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['H7', 'H8', 'H9', 'H10', 'H11', 'H12']]:
        pipette.mix(3, 100, well)
    pipette.drop_tip()
    
    # =========================================================================
    # PART 6: 10-minute pause for CaC2O4 precipitation
    # =========================================================================
    protocol.comment("=== PART 6: Pausing 10 min for CaC2O4 precipitation ===")
    protocol.delay(minutes=10)
    
    # =========================================================================
    # PART 7: Phase B Step 2 - NaOH additions (after oxalic acid pre-treatment)
    # =========================================================================
    protocol.comment("=== PART 7: Phase B Step 2 - NaOH additions ===")
    
    # NaOH 0.50M additions (F1-F3, F10-F12, G7-G9)
    pipette.transfer(30, reservoir1['A3'], [plate[w] for w in ['F1', 'F2', 'F3']], new_tip='once')
    pipette.transfer(30, reservoir1['A3'], [plate[w] for w in ['F10', 'F11', 'F12']], new_tip='once')
    pipette.transfer(30, reservoir1['A3'], [plate[w] for w in ['G7', 'G8', 'G9']], new_tip='once')
    
    # NaOH 0.75M additions (F4-F6, G1-G3, G10-G12)
    pipette.transfer(30, reservoir1['A4'], [plate[w] for w in ['F4', 'F5', 'F6']], new_tip='once')
    pipette.transfer(30, reservoir1['A4'], [plate[w] for w in ['G1', 'G2', 'G3']], new_tip='once')
    pipette.transfer(30, reservoir1['A4'], [plate[w] for w in ['G10', 'G11', 'G12']], new_tip='once')
    
    # NaOH 1.00M additions (F7-F9, G4-G6, H1-H3)
    pipette.transfer(30, reservoir1['A5'], [plate[w] for w in ['F7', 'F8', 'F9']], new_tip='once')
    pipette.transfer(30, reservoir1['A5'], [plate[w] for w in ['G4', 'G5', 'G6']], new_tip='once')
    pipette.transfer(30, reservoir1['A5'], [plate[w] for w in ['H1', 'H2', 'H3']], new_tip='once')
    
    # DI water for controls (H4-H12)
    pipette.transfer(30, reservoir2['A4'], [plate[w] for w in ['H4', 'H5', 'H6']], new_tip='once')
    pipette.transfer(30, reservoir2['A4'], [plate[w] for w in ['H7', 'H8', 'H9']], new_tip='once')
    pipette.transfer(30, reservoir2['A4'], [plate[w] for w in ['H10', 'H11', 'H12']], new_tip='once')
    
    # =========================================================================
    # PART 8: Phase B Step 2 mixing (after NaOH addition)
    # =========================================================================
    protocol.comment("=== PART 8: Phase B Step 2 mixing ===")
    
    # Group 1: NaOH 0.50M (F1-F3, F10-F12, G7-G9)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['F1', 'F2', 'F3', 'F10', 'F11', 'F12', 'G7', 'G8', 'G9']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 2: NaOH 0.75M (F4-F6, G1-G3, G10-G12)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['F4', 'F5', 'F6', 'G1', 'G2', 'G3', 'G10', 'G11', 'G12']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 3: NaOH 1.00M (F7-F9, G4-G6, H1-H3)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['F7', 'F8', 'F9', 'G4', 'G5', 'G6', 'H1', 'H2', 'H3']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    # Group 4: Controls/DI (H4-H12)
    pipette.pick_up_tip()
    for well in [plate[w] for w in ['H4', 'H5', 'H6', 'H7', 'H8', 'H9', 'H10', 'H11', 'H12']]:
        pipette.mix(5, 100, well)
    pipette.drop_tip()
    
    protocol.comment("=== Protocol complete. Seal plate, incubate 60 min RT, centrifuge 2000xg 10 min ===")
