from opentrons import protocol_api

metadata = {
    'protocolName': 'Reductive Acid Leaching Control - Un-roasted NdFeB',
    'author': 'OpentronsAI (corrected)',
    'description': '''Control experiment: 0.5M H2SO4 + 0.1M ascorbic acid on un-roasted NdFeB 
    powder. Tests whether mild reducing conditions selectively dissolve Fe while leaving 
    REEs behind. 6 temperatures (rows B-G) x 10 replicates (columns 2-11) = 60 wells.
    Complementary to the oxidative roast + selective leach main experiment.''',
    'source': 'OpentronsAI'
}

requirements = {
    'robotType': 'OT-2',
    'apiLevel': '2.15'
}

def run(protocol: protocol_api.ProtocolContext):

    # ========== LABWARE ==========
    reaction_plate = protocol.load_labware('nest_96_wellplate_2ml_deep', 1,
        'Reaction Plate (10mg un-roasted NdFeB in B2-G11)')
    acid_reservoir = protocol.load_labware('nest_12_reservoir_15ml', 2,
        'Reductive Acid (A1 = 0.5M H2SO4 + 0.1M ascorbic acid)')
    reagent_reservoir = protocol.load_labware('nest_12_reservoir_15ml', 3,
        'Reagents (A1=DI water, A2=2% HNO3, A3=ICP-MS diluent)')
    collection_plate = protocol.load_labware('nest_96_wellplate_2ml_deep', 4,
        'Collection Plate')
    icpms_plate = protocol.load_labware('nest_96_wellplate_2ml_deep', 9,
        'ICP-MS Plate (1:10 dilution)')

    # Tip racks
    tips_1000 = protocol.load_labware('opentrons_96_tiprack_1000ul', 5)
    tips_300_1 = protocol.load_labware('opentrons_96_tiprack_300ul', 7)
    tips_300_2 = protocol.load_labware('opentrons_96_tiprack_300ul', 8)
    tips_300_3 = protocol.load_labware('opentrons_96_tiprack_300ul', 10)

    # ========== PIPETTES ==========
    p1000 = protocol.load_instrument('p1000_single_gen2', 'left', tip_racks=[tips_1000])
    p300 = protocol.load_instrument('p300_single_gen2', 'right',
        tip_racks=[tips_300_1, tips_300_2, tips_300_3])

    # ========== WELL DEFINITIONS ==========
    # Guard wells (full perimeter)
    guard_wells = []
    guard_wells.extend(reaction_plate.rows()[0])     # Row A, all columns
    guard_wells.extend(reaction_plate.rows()[7])     # Row H, all columns
    for row_idx in range(1, 7):                       # Column 1, rows B-G
        guard_wells.append(reaction_plate.rows()[row_idx][0])
    for row_idx in range(1, 7):                       # Column 12, rows B-G
        guard_wells.append(reaction_plate.rows()[row_idx][11])

    # Active wells B2-G11: 6 temps x 10 replicates = 60 wells
    # Row B=25C, C=40C, D=55C, E=65C, F=75C, G=85C
    active_wells = []
    collection_wells = []
    icpms_wells = []
    for row_idx in range(1, 7):
        for col_idx in range(1, 11):
            active_wells.append(reaction_plate.rows()[row_idx][col_idx])
            collection_wells.append(collection_plate.rows()[row_idx][col_idx])
            icpms_wells.append(icpms_plate.rows()[row_idx][col_idx])

    # ========== STEP 1: GUARD WELLS ==========
    protocol.comment('=== Step 1: Filling guard wells with 1500 uL DI water ===')
    p1000.pick_up_tip()
    for well in guard_wells:
        p1000.aspirate(1000, reagent_reservoir['A1'])
        p1000.dispense(1000, well)
        p1000.aspirate(500, reagent_reservoir['A1'])
        p1000.dispense(500, well)
    p1000.drop_tip()

    # ========== STEP 2: ACID DISPENSING ==========
    # Same reductive solution to ALL 60 wells (single tip — no contamination risk)
    protocol.comment('=== Step 2: Dispensing 1000 uL reductive acid to all active wells ===')
    protocol.comment('Solution: 0.5M H2SO4 + 0.1M ascorbic acid')
    p1000.pick_up_tip()
    for well in active_wells:
        p1000.aspirate(1000, acid_reservoir['A1'])
        p1000.dispense(1000, well)
    p1000.drop_tip()

    # ========== STEP 3: INCUBATION PAUSE ==========
    protocol.comment('=== Step 3: Incubation pause ===')
    protocol.pause(
        'PROTOCOL PAUSED - Manual steps required:\n'
        '\n'
        'PRE-INCUBATION:\n'
        '1. Remove reaction plate from slot 1\n'
        '2. Centrifuge at 500xg for 30s (ensure powder-acid contact)\n'
        '3. Weigh plate (pre-incubation mass)\n'
        '4. Seal with acid-resistant aluminum foil\n'
        '\n'
        'INCUBATION:\n'
        '5. Incubate 4 hours at row temperatures with 300 rpm shaking:\n'
        '   Row B=25C, Row C=40C, Row D=55C, Row E=65C, Row F=75C, Row G=85C\n'
        '\n'
        'POST-INCUBATION:\n'
        '6. Cool to RT, weigh plate (verify <5% evaporative loss)\n'
        '7. Centrifuge at 3000xg for 10 min\n'
        '8. Remove seal, place plate back in slot 1\n'
        '9. Press RESUME'
    )

    # ========== STEP 4: SUPERNATANT TRANSFER ==========
    protocol.comment('=== Step 4: Transferring 500 uL supernatant (slow, 2mm above pellet) ===')
    p300.flow_rate.aspirate = 50  # Slow aspiration to avoid pellet
    for source, dest in zip(active_wells, collection_wells):
        p300.pick_up_tip()
        p300.aspirate(500, source.bottom(z=2))
        p300.dispense(500, dest)
        p300.drop_tip()
    p300.flow_rate.aspirate = 92.86  # Reset to default

    # ========== STEP 5: STABILIZATION ==========
    protocol.comment('=== Step 5: Adding 500 uL 2% HNO3 stabilizer (single tip) ===')
    p300.pick_up_tip()
    for dest in collection_wells:
        p300.aspirate(500, reagent_reservoir['A2'])
        p300.dispense(500, dest)
    p300.drop_tip()

    # ========== STEP 6a: ICP-MS DILUENT PRE-FILL ==========
    protocol.comment('=== Step 6a: Pre-filling ICP-MS plate with 900 uL diluent ===')
    p1000.pick_up_tip()
    for dest in icpms_wells:
        p1000.aspirate(900, reagent_reservoir['A3'])
        p1000.dispense(900, dest)
    p1000.drop_tip()

    # ========== STEP 6b: SAMPLE TRANSFER + MIX ==========
    protocol.comment('=== Step 6b: Transferring 100 uL sample for 1:10 dilution ===')
    for source, dest in zip(collection_wells, icpms_wells):
        p300.pick_up_tip()
        p300.aspirate(100, source)
        p300.dispense(100, dest)
        p300.mix(3, 200, dest)
        p300.drop_tip()

    # ========== COMPLETE ==========
    protocol.comment('=== PROTOCOL COMPLETE ===')
    protocol.comment('Slot 1: Reaction plate with residue pellets')
    protocol.comment('Slot 4: Collection plate (500uL leachate + 500uL HNO3)')
    protocol.comment('Slot 9: ICP-MS plate (1:10 dilution with In/Rh standards)')
    protocol.comment('Compare Fe vs REE selectivity against oxidative roast plates')
