from opentrons import protocol_api

metadata = {
    'protocolName': 'NdFeB Acid Leaching Screen - High Throughput',
    'author': 'OpentronsAI',
    'description': '''High-throughput acid leaching screen for NdFeB powder.
                      Phase 1: Acid dispensing (P1000 single-channel)
                      Phase 2: Supernatant transfer (P300 8-channel, 200 uL)
                      Phase 3: ICP-MS dilution (P300 8-channel)''',
    'source': 'OpentronsAI'
}

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

def run(protocol: protocol_api.ProtocolContext):
    
    # ===================== LABWARE SETUP =====================
    protocol.comment("=== Loading Labware ===")
    
    # Tip racks
    tips_1000_1 = protocol.load_labware('opentrons_96_tiprack_1000ul', 1)
    tips_1000_2 = protocol.load_labware('opentrons_96_tiprack_1000ul', 4)
    tips_300_1 = protocol.load_labware('opentrons_96_tiprack_300ul', 3)
    tips_300_2 = protocol.load_labware('opentrons_96_tiprack_300ul', 6)
    
    # Reservoir with acid stocks and diluent
    reservoir = protocol.load_labware('nest_12_reservoir_15ml', 2)
    
    # Deep well plates
    leaching_plate = protocol.load_labware('nest_96_wellplate_2ml_deep', 5)
    
    # Pipettes
    p1000 = protocol.load_instrument(
        'p1000_single_gen2', 
        'left', 
        tip_racks=[tips_1000_1, tips_1000_2]
    )
    
    p300_multi = protocol.load_instrument(
        'p300_multi_gen2', 
        'right', 
        tip_racks=[tips_300_1, tips_300_2]
    )
    
    # ===================== RESERVOIR LAYOUT =====================
    # A1: 0.10 M acid       A4: 1.00 M acid       A7: 1.00 M acid (row G blank)
    # A2: 0.25 M acid       A5: 2.00 M acid       A8: Ultrapure water (row H blank)
    # A3: 0.50 M acid       A6: 4.00 M acid       A9: 2% HNO3 (ICP-MS dilution)
    
    acid_sources = {
        'A': reservoir['A1'],  # 0.10 M
        'B': reservoir['A2'],  # 0.25 M
        'C': reservoir['A3'],  # 0.50 M
        'D': reservoir['A4'],  # 1.00 M
        'E': reservoir['A5'],  # 2.00 M
        'F': reservoir['A6'],  # 4.00 M
        'G': reservoir['A7'],  # 1.00 M blank
        'H': reservoir['A8']   # Water blank
    }
    
    # ===================== PHASE 1: ACID DISPENSING =====================
    protocol.comment("=" * 50)
    protocol.comment("PHASE 1: Dispensing acids to leaching plate")
    protocol.comment("Using P1000 single-channel, one tip per row")
    protocol.comment("=" * 50)
    
    # Dispense 1000 uL acid to each row
    for row_letter in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']:
        protocol.comment(f"Dispensing acid to row {row_letter}")
        
        # Get the row wells
        row_wells = leaching_plate.rows_by_name()[row_letter]
        acid_source = acid_sources[row_letter]
        
        # Pick up one tip for the entire row
        p1000.pick_up_tip()
        
        # Dispense to each well in the row
        for well in row_wells:
            p1000.aspirate(1000, acid_source, rate=0.5)
            protocol.delay(seconds=1)
            p1000.dispense(1000, well, rate=0.5)
            p1000.touch_tip(well)
            p1000.blow_out(well)
            protocol.delay(seconds=1)
        
        # Drop tip after finishing the row
        p1000.drop_tip()
        
        protocol.comment(f"Row {row_letter} complete")
    
    protocol.comment("=" * 50)
    protocol.comment("Phase 1 Complete: All acids dispensed")
    protocol.comment("Tips used in Phase 1: 8 tips (P1000)")
    protocol.comment("=" * 50)
    
    # ===================== PAUSE FOR SAMPLE ADDITION =====================
    protocol.pause(
        "===== PAUSE FOR SAMPLE PREPARATION =====\n"
        "\n"
        "Please perform the following steps:\n"
        "\n"
        "1. Remove the leaching plate from slot 5\n"
        "2. Add 10 mg NdFeB powder to wells in rows A-F and row H\n"
        "   (Do NOT add powder to row G - acid blank)\n"
        "3. Seal the plate with chemically resistant adhesive film\n"
        "4. Incubate at target temperature with shaking at 500 rpm for 120 min\n"
        "5. Centrifuge the sealed plate at 3000g for 10 min\n"
        "6. Remove the seal carefully\n"
        "7. Place the leaching plate back in slot 5\n"
        "8. Load empty deep well plates:\n"
        "   - Slot 8: Supernatant plate (nest_96_wellplate_2ml_deep)\n"
        "   - Slot 11: Dilution plate (nest_96_wellplate_2ml_deep)\n"
        "9. Ensure reservoir A9 contains 2% HNO3 for dilution\n"
        "\n"
        "Press RESUME when ready to continue with Phase 2."
    )
    
    # Load plates for Phase 2 (user places them during pause)
    supernatant_plate = protocol.load_labware('nest_96_wellplate_2ml_deep', 8)
    dilution_plate = protocol.load_labware('nest_96_wellplate_2ml_deep', 11)
    
    # ===================== PHASE 2: SUPERNATANT TRANSFER =====================
    protocol.comment("=" * 50)
    protocol.comment("PHASE 2: Transferring supernatant (200 uL)")
    protocol.comment("Using P300 8-channel, aspirating 2 mm above pellet")
    protocol.comment("=" * 50)
    
    # Transfer 200 uL supernatant column by column
    for col_idx in range(12):
        protocol.comment(f"Transferring column {col_idx + 1}")
        
        p300_multi.pick_up_tip()
        
        # Aspirate from 2 mm above bottom to avoid pellet
        p300_multi.aspirate(
            200, 
            leaching_plate.columns()[col_idx][0].bottom(2), 
            rate=0.3
        )
        
        # Dispense to supernatant plate
        p300_multi.dispense(200, supernatant_plate.columns()[col_idx][0])
        
        # Fresh tips per column
        p300_multi.drop_tip()
    
    protocol.comment("=" * 50)
    protocol.comment("Phase 2 Complete: All supernatant transferred")
    protocol.comment("Tips used in Phase 2: 96 tips (P300)")
    protocol.comment("=" * 50)
    
    # ===================== PAUSE BETWEEN PHASES =====================
    protocol.pause(
        "===== PAUSE BETWEEN PHASE 2 AND PHASE 3 =====\n"
        "\n"
        "Supernatant transfer complete.\n"
        "Verify that:\n"
        "- Supernatant plate in slot 8 contains transferred samples\n"
        "- Dilution plate in slot 11 is empty and ready\n"
        "- Reservoir A9 contains sufficient 2% HNO3\n"
        "\n"
        "Press RESUME to continue with ICP-MS dilution (Phase 3)."
    )
    
    # ===================== PHASE 3: ICP-MS DILUTION =====================
    protocol.comment("=" * 50)
    protocol.comment("PHASE 3: ICP-MS dilution (2x)")
    protocol.comment("100 uL supernatant + 100 uL 2% HNO3 = 200 uL total")
    protocol.comment("=" * 50)
    
    nitric_acid = reservoir['A9']
    
    for col_idx in range(12):
        protocol.comment(f"Diluting column {col_idx + 1}")
        
        p300_multi.pick_up_tip()
        
        # Transfer 100 uL supernatant to dilution plate
        p300_multi.aspirate(100, supernatant_plate.columns()[col_idx][0])
        p300_multi.dispense(100, dilution_plate.columns()[col_idx][0])
        
        # Add 100 uL of 2% HNO3
        p300_multi.aspirate(100, nitric_acid)
        p300_multi.dispense(100, dilution_plate.columns()[col_idx][0])
        
        # Mix 3 times with 150 uL
        p300_multi.mix(3, 150, dilution_plate.columns()[col_idx][0])
        
        p300_multi.drop_tip()
    
    # ===================== PROTOCOL COMPLETE =====================
    protocol.comment("=" * 50)
    protocol.comment("PROTOCOL COMPLETE")
    protocol.comment("=" * 50)
    protocol.comment("")
    protocol.comment("Total tips used:")
    protocol.comment("  Phase 1 (P1000): 8 tips")
    protocol.comment("  Phase 2 (P300):  96 tips")
    protocol.comment("  Phase 3 (P300):  96 tips")
    protocol.comment("  TOTAL:           200 tips")
    protocol.comment("")
    protocol.comment("Next steps:")
    protocol.comment("1. Label dilution plate (slot 11): date, acid type, temp, 2x dilution")
    protocol.comment("2. Seal dilution plate and store at 4 C until ICP-MS")
    protocol.comment("3. Store supernatant plate (slot 8) as backup at 4 C")
    protocol.comment("4. For high-acid wells (rows D-F), prepare additional manual dilutions")
    protocol.comment("   if ICP-MS signals are saturated (10 uL + 190 uL 2% HNO3 = 20x)")
    protocol.comment("")
    protocol.comment("Plate layout reference:")
    protocol.comment("  Row A: 0.10 M acid       Row E: 2.00 M acid")
    protocol.comment("  Row B: 0.25 M acid       Row F: 4.00 M acid")
    protocol.comment("  Row C: 0.50 M acid       Row G: 1.00 M acid blank (no solid)")
    protocol.comment("  Row D: 1.00 M acid       Row H: Water blank (solid + water)")
