from opentrons import protocol_api

metadata = {
    'protocolName': 'NdFeB Selective Acid Leaching - Single Roasting Temperature',
    'author': 'OpentronsAI',
    'description': '''Two-stage selective acid leaching of oxidatively roasted NdFeB magnet waste powder.
                      This protocol handles ONE plate (one roasting temperature).
                      Run this protocol 4 times, changing the roasted powder batch each time.
                      
                      ROASTING TEMPERATURE FOR THIS RUN: [USER: EDIT THIS VALUE]
                      
                      Active wells (B2-G11) contain 10mg pre-weighed roasted NdFeB powder.
                      Acid concentration range: 0.05-3.00 M HCl (10 concentrations).
                      Temperature range: 25-85C (6 temperatures, row-wise).''',
    'source': 'OpentronsAI'
}

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

def run(protocol: protocol_api.ProtocolContext):
    
    # ===================== LABWARE SETUP =====================
    
    # Slot 1: Reaction plate (already contains 10mg roasted powder in active wells B2-G11)
    reaction_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep',
        location=1,
        label='Reaction Plate (10mg NdFeB powder in B2-G11)'
    )
    
    # Slot 2: Acid stock reservoir (10 HCl concentrations in A1-A10)
    acid_reservoir = protocol.load_labware(
        'nest_12_reservoir_15ml',
        location=2,
        label='Acid Stocks (0.05-3.00 M HCl)'
    )
    
    # Slot 3: Reagent reservoir (DI water, stabilization acid, ICP-MS diluent)
    reagent_reservoir = protocol.load_labware(
        'nest_12_reservoir_15ml',
        location=3,
        label='Reagents (DI water, 2% HNO3, ICP-MS diluent)'
    )
    
    # Slot 4: Collection plate for supernatant
    collection_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep',
        location=4,
        label='Collection Plate (for supernatant)'
    )
    
    # Slot 5-6: 1000 uL tip racks for P1000
    tips_1000_1 = protocol.load_labware('opentrons_96_tiprack_1000ul', 5)
    tips_1000_2 = protocol.load_labware('opentrons_96_tiprack_1000ul', 6)
    
    # Slot 7-8: 300 uL tip racks for P300
    tips_300_1 = protocol.load_labware('opentrons_96_tiprack_300ul', 7)
    tips_300_2 = protocol.load_labware('opentrons_96_tiprack_300ul', 8)
    
    # Slot 9: ICP-MS analysis plate (for 1:10 diluted samples)
    icpms_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep',
        location=9,
        label='ICP-MS Plate (1:10 dilution)'
    )
    
    # Slot 10: Additional 300 uL tip rack
    tips_300_3 = protocol.load_labware('opentrons_96_tiprack_300ul', 10)
    
    # ===================== PIPETTE SETUP =====================
    
    # Left mount: P1000 single-channel GEN2
    p1000 = protocol.load_instrument(
        'p1000_single_gen2',
        mount='left',
        tip_racks=[tips_1000_1, tips_1000_2]
    )
    
    # Right mount: P300 single-channel GEN2
    p300 = protocol.load_instrument(
        'p300_single_gen2',
        mount='right',
        tip_racks=[tips_300_1, tips_300_2, tips_300_3]
    )
    
    # ===================== WELL DEFINITIONS =====================
    
    # Guard wells (perimeter)
    guard_wells = []
    # Row A (all columns)
    guard_wells.extend(reaction_plate.rows()[0])
    # Row H (all columns)
    guard_wells.extend(reaction_plate.rows()[7])
    # Column 1 (rows B-G)
    for row_idx in range(1, 7):  # rows B-G (indices 1-6)
        guard_wells.append(reaction_plate.rows()[row_idx][0])
    # Column 12 (rows B-G)
    for row_idx in range(1, 7):  # rows B-G (indices 1-6)
        guard_wells.append(reaction_plate.rows()[row_idx][11])
    
    # Active wells: B2-G11 (6 rows x 10 columns = 60 wells)
    # Row B=25C, C=40C, D=55C, E=65C, F=75C, G=85C
    # Columns 2-11 = 0.05, 0.10, 0.20, 0.30, 0.50, 0.75, 1.00, 1.50, 2.00, 3.00 M HCl
    active_wells = []
    for row_idx in range(1, 7):  # rows B-G (indices 1-6)
        for col_idx in range(1, 11):  # columns 2-11 (indices 1-10)
            active_wells.append(reaction_plate.rows()[row_idx][col_idx])
    
    # ===================== STEP 1: GUARD WELLS =====================
    protocol.comment('=== STEP 1: Filling guard wells with 1500 uL DI water ===')
    
    # DI water is in reagent_reservoir well A1
    di_water = reagent_reservoir.wells()[0]
    
    # Pick up one tip for all guard well transfers
    p1000.pick_up_tip()
    
    # Fill each guard well with 1500 uL (split into 1000 + 500 uL)
    for well in guard_wells:
        # First 1000 uL
        p1000.aspirate(1000, di_water)
        p1000.dispense(1000, well)
        # Second 500 uL
        p1000.aspirate(500, di_water)
        p1000.dispense(500, well)
    
    p1000.drop_tip()
    
    protocol.comment(f'Guard wells filled: {len(guard_wells)} wells with 1500 uL DI water each')
    
    # ===================== STEP 2: ACID DISPENSING =====================
    protocol.comment('=== STEP 2: Dispensing 1000 uL acid to active wells (low to high concentration) ===')
    
    # Acid concentrations in reservoir A1-A10: 0.05, 0.10, 0.20, 0.30, 0.50, 0.75, 1.00, 1.50, 2.00, 3.00 M
    acid_concentrations = [0.05, 0.10, 0.20, 0.30, 0.50, 0.75, 1.00, 1.50, 2.00, 3.00]
    
    # Dispense from lowest to highest concentration
    for acid_idx in range(10):  # 10 acid concentrations
        # Get acid source well (A1-A10 in acid_reservoir)
        acid_source = acid_reservoir.wells()[acid_idx]
        
        # Get destination column (columns 2-11, which are indices 1-10)
        dest_col_idx = acid_idx + 1
        
        # Get destination wells for this column (rows B-G only)
        dest_wells = []
        for row_idx in range(1, 7):  # rows B-G (indices 1-6)
            dest_wells.append(reaction_plate.rows()[row_idx][dest_col_idx])
        
        # Pick up fresh tip for this acid concentration
        p1000.pick_up_tip()
        
        protocol.comment(f'Dispensing {acid_concentrations[acid_idx]} M HCl to column {dest_col_idx + 1} (rows B-G)')
        
        # Dispense 1000 uL to each well in this column
        for well in dest_wells:
            p1000.aspirate(1000, acid_source)
            p1000.dispense(1000, well)
        
        p1000.drop_tip()
    
    protocol.comment('Acid dispensing complete: 60 active wells filled with 1000 uL acid each')
    
    # ===================== STEP 2b: BLANKS/CONTROLS NOTE =====================
    protocol.comment('NOTE: Prepare off-plate controls in 8-well strip tubes:')
    protocol.comment('  - 3 acid blanks (0.50M HCl, no solid, 55C, 4h)')
    protocol.comment('  - 3 positive controls (spiked REE/Fe standards)')
    protocol.comment('  - 3 un-roasted NdFeB comparators (0.50M HCl, 55C, 4h)')
    
    # ===================== 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 30 seconds (ensure powder-acid contact)\n'
        '3. Weigh the plate (record pre-incubation mass)\n'
        '4. Seal with acid-resistant aluminum foil\n'
        '\n'
        'INCUBATION:\n'
        '5. Incubate 4 hours at row-specific 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'
        '   (Use thermal gradient block or sequential temperature batches)\n'
        '\n'
        'POST-INCUBATION:\n'
        '6. Cool to room temperature (~5 min)\n'
        '7. Weigh the plate again (verify <5% evaporative loss)\n'
        '8. Centrifuge at 3000xg for 10 minutes (pellet residue)\n'
        '9. Carefully remove seal\n'
        '10. Place reaction plate back in slot 1\n'
        '\n'
        'Press RESUME to continue with supernatant transfer'
    )
    
    # ===================== STEP 4: SUPERNATANT TRANSFER =====================
    protocol.comment('=== STEP 4: Transferring 500 uL supernatant to collection plate ===')
    
    # Set slow aspiration rate (50 uL/s) to avoid disturbing pellet
    p300.flow_rate.aspirate = 50
    
    # Transfer supernatant from each active well
    for idx, source_well in enumerate(active_wells):
        # Get corresponding destination well in collection plate
        # Active wells are B2-G11, same positions in collection plate
        row_idx = (idx // 10) + 1  # rows B-G (indices 1-6)
        col_idx = (idx % 10) + 1   # columns 2-11 (indices 1-10)
        dest_well = collection_plate.rows()[row_idx][col_idx]
        
        # Pick up fresh tip for each well
        p300.pick_up_tip()
        
        # Aspirate 500 uL from 2 mm above bottom (avoid pellet)
        p300.aspirate(500, source_well.bottom(z=2))
        
        # Dispense into collection plate
        p300.dispense(500, dest_well)
        
        p300.drop_tip()
    
    # Reset flow rate to default
    p300.flow_rate.aspirate = 92.86  # default for P300 GEN2
    
    protocol.comment('Supernatant transfer complete: 60 wells transferred')
    
    # ===================== STEP 5: STABILIZATION =====================
    protocol.comment('=== STEP 5: Adding 500 uL 2% HNO3 stabilization acid ===')
    
    # 2% HNO3 is in reagent_reservoir well A2
    stabilization_acid = reagent_reservoir.wells()[1]
    
    # Pick up one tip for all stabilization additions
    p300.pick_up_tip()
    
    # Add 500 uL to each active well in collection plate
    for idx in range(60):
        row_idx = (idx // 10) + 1  # rows B-G (indices 1-6)
        col_idx = (idx % 10) + 1   # columns 2-11 (indices 1-10)
        dest_well = collection_plate.rows()[row_idx][col_idx]
        
        p300.aspirate(500, stabilization_acid)
        p300.dispense(500, dest_well)
    
    p300.drop_tip()
    
    protocol.comment('Stabilization complete: 60 wells now contain 1000 uL total (500 uL sample + 500 uL 2% HNO3)')
    
    # ===================== STEP 6a: ICP-MS DILUENT PRE-FILL =====================
    protocol.comment('=== STEP 6a: Pre-filling ICP-MS plate with 900 uL diluent ===')
    
    # ICP-MS diluent (2% HNO3 + 10ppb In + 10ppb Rh) is in reagent_reservoir well A3
    icpms_diluent = reagent_reservoir.wells()[2]
    
    # Pick up one tip for all diluent additions
    p1000.pick_up_tip()
    
    # Add 900 uL diluent to each active well in ICP-MS plate
    for idx in range(60):
        row_idx = (idx // 10) + 1  # rows B-G (indices 1-6)
        col_idx = (idx % 10) + 1   # columns 2-11 (indices 1-10)
        dest_well = icpms_plate.rows()[row_idx][col_idx]
        
        p1000.aspirate(900, icpms_diluent)
        p1000.dispense(900, dest_well)
    
    p1000.drop_tip()
    
    protocol.comment('ICP-MS plate pre-filled: 60 wells with 900 uL diluent each')
    
    # ===================== STEP 6b: SAMPLE DILUTION =====================
    protocol.comment('=== STEP 6b: Transferring 100 uL sample for 1:10 dilution ===')
    
    # Transfer 100 uL from collection plate to ICP-MS plate with mixing
    for idx in range(60):
        row_idx = (idx // 10) + 1  # rows B-G (indices 1-6)
        col_idx = (idx % 10) + 1   # columns 2-11 (indices 1-10)
        
        source_well = collection_plate.rows()[row_idx][col_idx]
        dest_well = icpms_plate.rows()[row_idx][col_idx]
        
        # Pick up fresh tip for each well
        p300.pick_up_tip()
        
        # Transfer 100 uL
        p300.aspirate(100, source_well)
        p300.dispense(100, dest_well)
        
        # Mix 3 times with 200 uL
        p300.mix(3, 200, dest_well)
        
        p300.drop_tip()
    
    protocol.comment('ICP-MS dilution complete: 60 wells with 1000 uL total (1:10 dilution)')
    
    # ===================== PROTOCOL SUMMARY =====================
    protocol.comment('=== PROTOCOL COMPLETE ===')
    protocol.comment('Plates ready for analysis:')
    protocol.comment('  Slot 1 - Reaction plate: leached residue pellets (can dry/weigh for mass balance)')
    protocol.comment('  Slot 4 - Collection plate: 1000 uL stabilized leachate per well')
    protocol.comment('  Slot 9 - ICP-MS plate: 1000 uL at 1:10 dilution with In/Rh standards')
    protocol.comment('Remember: Run off-plate blanks and controls through ICP-MS')
    protocol.comment('Repeat this protocol with next roasting temperature batch')
    protocol.comment(f'Tip usage: ~12 P1000 tips, ~181 P300 tips')
