from opentrons import protocol_api

metadata = {
    'protocolName': 'NdFeB Magnet Waste Acid Leaching - Plate 1 (HCl)',
    'author': 'OpentronsAI',
    'description': '''96-deep-well plate acid leaching protocol for rare earth element
    recovery from NdFeB magnet waste. Full-factorial broad screen: 10 HCl concentrations
    x 6 temperatures = 60 unique conditions (no replicates). Includes guard wells,
    supernatant transfer, HNO3 stabilization, and 1:10 ICP-MS dilution with In/Rh
    internal standards. Replication deferred to Bayesian optimization phase.''',
    'source': 'OpentronsAI'
}

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

def run(protocol: protocol_api.ProtocolContext):

    # =====================================================================
    # LABWARE SETUP
    # =====================================================================

    # Slot 1: Reaction plate (10mg NdFeB powder pre-weighed into active wells)
    reaction_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep', location='1',
        label='Reaction Plate - HCl Leaching'
    )

    # Slot 2: 12-well reservoir with 10 HCl concentrations (wells A1-A10)
    #   A1=0.10M, A2=0.25M, A3=0.50M, A4=0.75M, A5=1.00M,
    #   A6=1.50M, A7=2.00M, A8=3.00M, A9=4.00M, A10=6.00M
    acid_reservoir = protocol.load_labware(
        'nest_12_reservoir_15ml', location='2',
        label='HCl Stock Solutions (0.1-6.0 M)'
    )

    # Slot 3: 12-well reservoir
    #   A1 = DI water (for guard wells)
    #   A2 = 2% HNO3 (for stabilization)
    #   A3 = 2% HNO3 + 10 ppb In + 10 ppb Rh (ICP-MS diluent with internal standards)
    water_reservoir = protocol.load_labware(
        'nest_12_reservoir_15ml', location='3',
        label='DI Water (A1), 2% HNO3 (A2), ICP-MS Diluent (A3)'
    )

    # Slot 4: Collection plate for supernatant after centrifugation
    collection_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep', location='4',
        label='Collection Plate'
    )

    # Slot 9: ICP-MS analysis plate (1:10 diluted samples ready for measurement)
    icpms_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep', location='9',
        label='ICP-MS Analysis Plate (1:10 diluted)'
    )

    # Tip racks
    tiprack_1000_1 = protocol.load_labware('opentrons_96_tiprack_1000ul', '5')
    tiprack_1000_2 = protocol.load_labware('opentrons_96_tiprack_1000ul', '6')
    tiprack_300_1 = protocol.load_labware('opentrons_96_tiprack_300ul', '7')
    tiprack_300_2 = protocol.load_labware('opentrons_96_tiprack_300ul', '8')
    tiprack_300_3 = protocol.load_labware('opentrons_96_tiprack_300ul', '10')

    # =====================================================================
    # PIPETTE SETUP
    # =====================================================================

    p1000 = protocol.load_instrument(
        'p1000_single_gen2', mount='left',
        tip_racks=[tiprack_1000_1, tiprack_1000_2]
    )

    p300 = protocol.load_instrument(
        'p300_single_gen2', mount='right',
        tip_racks=[tiprack_300_1, tiprack_300_2, tiprack_300_3]
    )

    # =====================================================================
    # WELL DEFINITIONS
    # =====================================================================
    # Full-factorial broad screen (no replicates):
    #   60 unique conditions = 10 acid concentrations x 6 temperatures
    #   Guard wells: Row A (all), Row H (all), Col 1 (B-G), Col 12 (B-G)
    #   Active wells: Columns 2-11, Rows B-G
    #
    # Temperature assignment per row:
    #   Row B=25C, Row C=40C, Row D=55C, Row E=65C, Row F=75C, Row G=85C
    #
    # Acid concentration per column:
    #   Col2=0.10M, Col3=0.25M, Col4=0.50M, Col5=0.75M, Col6=1.00M,
    #   Col7=1.50M, Col8=2.00M, Col9=3.00M, Col10=4.00M, Col11=6.00M
    #
    # Blanks/controls: Handled off-plate (see Step 2b comments)

    # Build guard well list (perimeter wells for evaporation barrier)
    guard_wells = []
    guard_wells.extend(reaction_plate.rows()[0])     # Row A, all 12 columns
    guard_wells.extend(reaction_plate.rows()[7])     # Row H, all 12 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])

    # Build active well lists (matched across reaction, collection, and ICP-MS plates)
    active_wells = []
    collection_wells = []
    icpms_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])
            collection_wells.append(collection_plate.rows()[row_idx][col_idx])
            icpms_wells.append(icpms_plate.rows()[row_idx][col_idx])

    # =====================================================================
    # STEP 1: FILL GUARD WELLS WITH DI WATER (evaporation barrier)
    # =====================================================================
    # Purpose: Prevent edge effects and evaporative losses in active wells
    # Volume: 1500 uL per guard well
    # Tip usage: Single tip (DI water only, no contamination risk)

    protocol.comment("=== STEP 1: Filling guard wells with 1500 uL DI water ===")
    protocol.comment(f"Total guard wells: {len(guard_wells)}")

    p1000.pick_up_tip()
    for well in guard_wells:
        # Split 1500 uL into two transfers (P1000 max = 1000 uL)
        p1000.aspirate(1000, water_reservoir['A1'])
        p1000.dispense(1000, well)
        p1000.aspirate(500, water_reservoir['A1'])
        p1000.dispense(500, well)
    p1000.drop_tip()

    protocol.comment("Guard wells filled successfully")

    # =====================================================================
    # STEP 2: DISPENSE ACID SOLUTIONS TO ACTIVE WELLS
    # =====================================================================
    # Purpose: Deliver precise acid volumes for leaching screen
    # Volume: 1000 uL per well (gives 1:100 solid:liquid ratio with 10mg powder)
    # Tip usage: Fresh tip per acid concentration (10 tips total)
    # Order: Lowest to highest concentration to minimize carryover risk

    protocol.comment("=== STEP 2: Dispensing HCl to active wells (low to high conc) ===")

    # Mapping: reservoir well -> plate column index -> concentration label
    acid_mapping = [
        ('A1',  1,  '0.10 M HCl'),   # Reservoir A1  -> Plate col 2
        ('A2',  2,  '0.25 M HCl'),   # Reservoir A2  -> Plate col 3
        ('A3',  3,  '0.50 M HCl'),   # Reservoir A3  -> Plate col 4
        ('A4',  4,  '0.75 M HCl'),   # Reservoir A4  -> Plate col 5
        ('A5',  5,  '1.00 M HCl'),   # Reservoir A5  -> Plate col 6
        ('A6',  6,  '1.50 M HCl'),   # Reservoir A6  -> Plate col 7
        ('A7',  7,  '2.00 M HCl'),   # Reservoir A7  -> Plate col 8
        ('A8',  8,  '3.00 M HCl'),   # Reservoir A8  -> Plate col 9
        ('A9',  9,  '4.00 M HCl'),   # Reservoir A9  -> Plate col 10
        ('A10', 10, '6.00 M HCl'),   # Reservoir A10 -> Plate col 11
    ]

    for reservoir_well, col_idx, conc_label in acid_mapping:
        protocol.comment(f"Dispensing {conc_label} -> plate column {col_idx + 1}")

        # Get the 6 wells in this column (rows B-G)
        column_wells = [reaction_plate.rows()[r][col_idx] for r in range(1, 7)]

        p1000.pick_up_tip()
        for well in column_wells:
            p1000.aspirate(1000, acid_reservoir[reservoir_well])
            p1000.dispense(1000, well)
        p1000.drop_tip()

    protocol.comment("Acid dispensing complete - all 60 active wells loaded")

    # =====================================================================
    # STEP 2b: BLANKS AND CONTROLS (off-plate, manual)
    # =====================================================================
    # Since all 60 active wells are used for the full-factorial screen,
    # blanks and controls are prepared separately:
    #   - 3 blank wells: acid only (no powder) in a separate strip tube
    #     at representative concentrations (e.g., 0.50M, 2.00M, 6.00M HCl)
    #   - 3 positive controls: spiked REE standard solutions
    #     (mixed Nd/Dy/Pr/Fe/Co at known concentrations in 2% HNO3)
    # These are processed through ICP-MS alongside plate samples.

    protocol.comment("NOTE: Prepare 3 blank tubes (acid only) and 3 positive controls")
    protocol.comment("(spiked REE standards) off-plate for ICP-MS QC")

    # =====================================================================
    # STEP 3: PAUSE FOR CENTRIFUGE, SEALING, INCUBATION, AND CENTRIFUGATION
    # =====================================================================

    protocol.comment("=== STEP 3: Pausing for off-deck processing ===")
    protocol.pause(
        "PROTOCOL PAUSED — Manual steps required:\n"
        "\n"
        "PRE-INCUBATION:\n"
        "1. Remove reaction plate from slot 1\n"
        "2. Centrifuge briefly at 500 x g for 30 seconds to ensure\n"
        "   all powder contacts acid and no droplets on well walls\n"
        "3. Weigh the plate and record mass (pre-incubation weight)\n"
        "4. Seal with chemically resistant aluminum foil\n"
        "\n"
        "INCUBATION:\n"
        "5. Incubate 4 hours at row-specific temperatures:\n"
        "   Row B=25C, Row C=40C, Row D=55C, Row E=65C, Row F=75C, Row G=85C\n"
        "   with 300 rpm orbital shaking\n"
        "   (Use thermal gradient block or sequential temperature batches)\n"
        "\n"
        "POST-INCUBATION:\n"
        "6. Cool to room temperature (~5 min on bench)\n"
        "7. Weigh the plate again (post-incubation weight)\n"
        "   -> Verify evaporative loss < 5% of pre-incubation weight\n"
        "8. Centrifuge at 3000 x g for 10 minutes to pellet residue\n"
        "9. Carefully remove seal\n"
        "10. Place reaction plate back in slot 1\n"
        "11. Press RESUME to continue supernatant transfer"
    )

    # =====================================================================
    # STEP 4: SUPERNATANT TRANSFER TO COLLECTION PLATE
    # =====================================================================
    # Purpose: Separate leachate from undissolved residue
    # Volume: 500 uL from each active well
    # Aspiration: Slow (50 uL/s) from 2mm above bottom to avoid pellet
    # Tip usage: Fresh tip per well (60 tips) to prevent cross-contamination

    protocol.comment("=== STEP 4: Transferring 500 uL supernatant to collection plate ===")
    protocol.comment("Slow aspiration (50 uL/s) at 2mm above bottom to avoid pellet")

    # Set slow aspiration rate
    p300.flow_rate.aspirate = 50  # uL/s (default is ~92 uL/s)

    for source_well, dest_well in zip(active_wells, collection_wells):
        p300.pick_up_tip()
        p300.aspirate(500, source_well.bottom(z=2))  # 2mm above bottom
        p300.dispense(500, dest_well)
        p300.drop_tip()

    # Reset to default flow rate
    p300.flow_rate.aspirate = 92.86

    protocol.comment(f"Supernatant transferred from {len(active_wells)} wells")

    # =====================================================================
    # STEP 5: HNO3 STABILIZATION OF COLLECTION PLATE
    # =====================================================================
    # Purpose: Acidify samples with 2% HNO3 to prevent metal hydrolysis
    #          and precipitation during storage
    # Volume: 500 uL per well (sample is now 500 uL leachate + 500 uL HNO3)
    # Tip usage: Single tip (same solution to all wells)

    protocol.comment("=== STEP 5: Adding 500 uL 2% HNO3 stabilizer to collection plate ===")

    p300.pick_up_tip()
    for dest_well in collection_wells:
        p300.aspirate(500, water_reservoir['A2'])   # 2% HNO3 in reservoir A2
        p300.dispense(500, dest_well)
    p300.drop_tip()

    protocol.comment("Stabilization complete — collection plate has 1000 uL per well")

    # =====================================================================
    # STEP 6: ICP-MS SAMPLE DILUTION (1:10)
    # =====================================================================
    # Purpose: Dilute stabilized samples to ICP-MS calibration range
    #          and add In/Rh internal standards for drift correction
    # Method: Transfer 100 uL from collection plate + 900 uL ICP-MS diluent
    #         (2% HNO3 containing 10 ppb In and 10 ppb Rh)
    # Final volume: 1000 uL at 1:10 dilution in ICP-MS plate
    # Tip usage: Fresh tip per well for transfer, single tip for diluent

    protocol.comment("=== STEP 6: ICP-MS 1:10 dilution with In/Rh internal standards ===")

    # First, pre-fill ICP-MS plate with 900 uL diluent (single tip)
    protocol.comment("Pre-filling ICP-MS plate with 900 uL diluent (2% HNO3 + In/Rh)")
    p1000.pick_up_tip()
    for dest_well in icpms_wells:
        p1000.aspirate(900, water_reservoir['A3'])  # ICP-MS diluent in reservoir A3
        p1000.dispense(900, dest_well)
    p1000.drop_tip()

    # Then transfer 100 uL of stabilized sample (fresh tip per well)
    protocol.comment("Transferring 100 uL stabilized sample to ICP-MS plate")
    for source_well, dest_well in zip(collection_wells, icpms_wells):
        p300.pick_up_tip()
        p300.aspirate(100, source_well)
        p300.dispense(100, dest_well)
        p300.mix(3, 200, dest_well)  # Mix 3x to homogenize
        p300.drop_tip()

    protocol.comment("ICP-MS dilution complete — plate ready for measurement")

    # =====================================================================
    # PROTOCOL COMPLETE
    # =====================================================================
    protocol.comment("=== PROTOCOL COMPLETE ===")
    protocol.comment("Summary:")
    protocol.comment("  Reaction plate (slot 1): 60 wells with leached residue")
    protocol.comment("  Collection plate (slot 4): 60 wells, 1000 uL stabilized leachate")
    protocol.comment("  ICP-MS plate (slot 9): 60 wells, 1000 uL at 1:10 dilution")
    protocol.comment("  Remember to run 3 blanks and 3 positive controls through ICP-MS")
    protocol.comment("")
    protocol.comment("Tip usage: ~11 P1000 tips, ~121 P300 tips + ~60 P300 for ICP-MS dilution")
    protocol.comment("Measure Nd, Dy, Pr, Fe, Co, B, Cu, Al, Ni by ICP-MS")
    protocol.comment("Calculate % recovery per element per condition")
