from opentrons import protocol_api

metadata = {
    'protocolName': 'NdFeB Acid Leaching Screen - Iteration 2: Oxidative Roasting + Selective Leach',
    'author': 'OpentronsAI',
    'description': '''Iteration 2: Oxidatively roasted NdFeB powder selective acid leaching.
                      8 roasting conditions (4 temps x 2 pretreatments) x 6 lixiviants x duplicates.
                      Phase 1: Acid dispensing (P1000 single-channel)
                      Phase 2: Supernatant transfer (P300 8-channel, 200 uL)
                      Phase 3: ICP-MS 10x dilution (P300 8-channel, 20+180 uL)''',
    'source': 'OpentronsAI'
}

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

def run(protocol: protocol_api.ProtocolContext):

    # ===================== LABWARE SETUP =====================
    protocol.comment("=" * 50)
    protocol.comment("ITERATION 2: Roasted NdFeB Selective Leaching Screen")
    protocol.comment("=" * 50)

    # 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 acids and diluent
    reservoir = protocol.load_labware('nest_12_reservoir_15ml', 2)

    # Plates
    experimental_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep', 5, 'Experimental Leaching Plate'
    )
    control_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep', 11, 'Control/Blank Plate'
    )

    # 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 MAPPING =====================
    # A1: 0.05 M H2SO4       A5: 1.00 M H2SO4
    # A2: 0.10 M H2SO4       A6: 1.00 M Acetic acid
    # A3: 0.25 M H2SO4       A7: Ultrapure water (solid blanks)
    # A4: 0.50 M H2SO4       A8: 2% HNO3 (ICP-MS dilution)

    acid_005M = reservoir['A1']
    acid_010M = reservoir['A2']
    acid_025M = reservoir['A3']
    acid_050M = reservoir['A4']
    acid_100M = reservoir['A5']
    acetic_acid = reservoir['A6']
    water = reservoir['A7']
    hno3_2pct = reservoir['A8']

    # ===================== PHASE 1: ACID DISPENSING =====================
    protocol.comment("=" * 50)
    protocol.comment("PHASE 1: ACID DISPENSING TO EXPERIMENTAL + CONTROL PLATES")
    protocol.comment("=" * 50)
    protocol.comment("")
    protocol.comment("Experimental plate layout:")
    protocol.comment("  Rows A-H = Roasting batches R1-R8")
    protocol.comment("    A=R1 (NaOH+400C)  B=R2 (NaOH+500C)  C=R3 (NaOH+600C)  D=R4 (NaOH+800C)")
    protocol.comment("    E=R5 (Direct+400C) F=R6 (Direct+500C) G=R7 (Direct+600C) H=R8 (Direct+800C)")
    protocol.comment("  Cols 1-2: 0.05M H2SO4   Cols 7-8:  0.50M H2SO4")
    protocol.comment("  Cols 3-4: 0.10M H2SO4   Cols 9-10: 1.00M H2SO4")
    protocol.comment("  Cols 5-6: 0.25M H2SO4   Cols 11-12: 1.00M Acetic acid")
    protocol.comment("")

    # --- Experimental plate: dispense acids (1 tip per acid type) ---
    # Each acid goes to 2 columns x 8 rows = 16 wells
    acid_mapping_exp = [
        (acid_005M,    "0.05M H2SO4",
         ['A1','B1','C1','D1','E1','F1','G1','H1',
          'A2','B2','C2','D2','E2','F2','G2','H2']),
        (acid_010M,    "0.10M H2SO4",
         ['A3','B3','C3','D3','E3','F3','G3','H3',
          'A4','B4','C4','D4','E4','F4','G4','H4']),
        (acid_025M,    "0.25M H2SO4",
         ['A5','B5','C5','D5','E5','F5','G5','H5',
          'A6','B6','C6','D6','E6','F6','G6','H6']),
        (acid_050M,    "0.50M H2SO4",
         ['A7','B7','C7','D7','E7','F7','G7','H7',
          'A8','B8','C8','D8','E8','F8','G8','H8']),
        (acid_100M,    "1.00M H2SO4",
         ['A9','B9','C9','D9','E9','F9','G9','H9',
          'A10','B10','C10','D10','E10','F10','G10','H10']),
        (acetic_acid,  "1.00M Acetic acid",
         ['A11','B11','C11','D11','E11','F11','G11','H11',
          'A12','B12','C12','D12','E12','F12','G12','H12']),
    ]

    protocol.comment("--- Dispensing to EXPERIMENTAL plate (slot 5) ---")
    for source, label, wells in acid_mapping_exp:
        protocol.comment(f"Dispensing {label} to {len(wells)} wells")
        p1000.pick_up_tip()
        for well_name in wells:
            p1000.aspirate(1000, source, rate=0.5)
            p1000.dispense(1000, experimental_plate[well_name], rate=0.5)
            p1000.touch_tip(experimental_plate[well_name])
            p1000.blow_out(experimental_plate[well_name])
            protocol.delay(seconds=1)
        p1000.drop_tip()
        protocol.comment(f"  {label} complete")

    # --- Control plate ---
    protocol.comment("")
    protocol.comment("--- Dispensing to CONTROL plate (slot 11) ---")

    # Row A: acid-only blanks (no solid), 2 wells per acid
    protocol.comment("Row A: acid-only blanks")
    control_row_a = [
        (acid_005M,   "0.05M H2SO4",   ['A1', 'A2']),
        (acid_010M,   "0.10M H2SO4",   ['A3', 'A4']),
        (acid_025M,   "0.25M H2SO4",   ['A5', 'A6']),
        (acid_050M,   "0.50M H2SO4",   ['A7', 'A8']),
        (acid_100M,   "1.00M H2SO4",   ['A9', 'A10']),
        (acetic_acid, "1.00M Acetic",   ['A11', 'A12']),
    ]
    for source, label, wells in control_row_a:
        p1000.pick_up_tip()
        for well_name in wells:
            p1000.aspirate(1000, source, rate=0.5)
            p1000.dispense(1000, control_plate[well_name], rate=0.5)
            p1000.touch_tip(control_plate[well_name])
            p1000.blow_out(control_plate[well_name])
            protocol.delay(seconds=1)
        p1000.drop_tip()

    # Row B: solid-only blanks (water + roasted powder), B1-B8
    protocol.comment("Row B: solid-only blanks (water + roasted powder)")
    p1000.pick_up_tip()
    for well_name in ['B1', 'B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8']:
        p1000.aspirate(1000, water, rate=0.5)
        p1000.dispense(1000, control_plate[well_name], rate=0.5)
        p1000.touch_tip(control_plate[well_name])
        p1000.blow_out(control_plate[well_name])
        protocol.delay(seconds=1)
    p1000.drop_tip()

    # Row C: unroasted metallic NdFeB reference, 2 wells per acid
    protocol.comment("Row C: unroasted metallic NdFeB reference")
    control_row_c = [
        (acid_005M,   "0.05M H2SO4",   ['C1', 'C2']),
        (acid_010M,   "0.10M H2SO4",   ['C3', 'C4']),
        (acid_025M,   "0.25M H2SO4",   ['C5', 'C6']),
        (acid_050M,   "0.50M H2SO4",   ['C7', 'C8']),
        (acid_100M,   "1.00M H2SO4",   ['C9', 'C10']),
        (acetic_acid, "1.00M Acetic",   ['C11', 'C12']),
    ]
    for source, label, wells in control_row_c:
        p1000.pick_up_tip()
        for well_name in wells:
            p1000.aspirate(1000, source, rate=0.5)
            p1000.dispense(1000, control_plate[well_name], rate=0.5)
            p1000.touch_tip(control_plate[well_name])
            p1000.blow_out(control_plate[well_name])
            protocol.delay(seconds=1)
        p1000.drop_tip()

    protocol.comment("")
    protocol.comment("=" * 50)
    protocol.comment("PHASE 1 COMPLETE")
    protocol.comment("=" * 50)
    protocol.comment("Tips used: ~19 P1000 tips")
    protocol.comment("  6 (exp acids) + 6 (ctrl row A) + 1 (ctrl row B water) + 6 (ctrl row C)")

    # ===================== PAUSE: INCUBATION =====================
    protocol.pause(
        "PHASE 1 COMPLETE. Solid has already been pre-dispensed manually.\n"
        "\n"
        "Please perform the following steps:\n"
        "1. Seal BOTH plates with chemically resistant adhesive film\n"
        "2. Place both plates on thermoshaker at 60 C, 500 rpm for 120 min\n"
        "3. After incubation, cool to 25 C (~10 min on bench)\n"
        "4. Centrifuge both plates at 3000g for 10 min\n"
        "5. Carefully remove seals\n"
        "6. Place experimental plate back in slot 5\n"
        "7. Place control plate back in slot 11\n"
        "8. Load empty deep-well plate in slot 8 (supernatant collection)\n"
        "9. Load empty deep-well plate in slot 9 (ICP-MS dilution)\n"
        "10. Ensure reservoir A8 contains 2% HNO3\n"
        "\n"
        "Press RESUME when ready for Phase 2."
    )

    # Load plates for Phase 2 (user places them during pause)
    supernatant_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep', 8, 'Supernatant Collection Plate'
    )
    dilution_plate = protocol.load_labware(
        'nest_96_wellplate_2ml_deep', 9, 'ICP-MS Dilution Plate'
    )

    # ===================== PHASE 2: SUPERNATANT TRANSFER =====================
    protocol.comment("")
    protocol.comment("=" * 50)
    protocol.comment("PHASE 2: SUPERNATANT TRANSFER (experimental plate only)")
    protocol.comment("Using P300 8-channel, 200 uL, aspirating 2mm above pellet")
    protocol.comment("=" * 50)

    for col_idx in range(12):
        protocol.comment(f"Transferring column {col_idx + 1} of 12")
        p300_multi.pick_up_tip()
        # Aspirate 200 uL from 2mm above bottom to avoid pellet
        p300_multi.aspirate(
            200,
            experimental_plate.columns()[col_idx][0].bottom(2),
            rate=0.3
        )
        p300_multi.dispense(
            200,
            supernatant_plate.columns()[col_idx][0]
        )
        p300_multi.drop_tip()

    protocol.comment("")
    protocol.comment("Experimental plate supernatant transfer complete")
    protocol.comment("Tips used in Phase 2: 96 P300 tips (12 columns x 8 channel)")
    protocol.comment("")
    protocol.comment("NOTE: Manually transfer control plate supernatant")
    protocol.comment("using a multichannel pipette for rows A, B, C")

    protocol.pause(
        "Supernatant transfer from experimental plate complete.\n"
        "\n"
        "Please:\n"
        "- Manually transfer control plate (slot 11) supernatant\n"
        "  using a multichannel pipette for rows A, B, C\n"
        "  into labeled tubes or a separate plate\n"
        "- Ensure ICP-MS dilution plate is in slot 9\n"
        "- Ensure 2% HNO3 in reservoir A8\n"
        "\n"
        "Press RESUME for Phase 3 (ICP-MS dilution)."
    )

    # ===================== PHASE 3: ICP-MS DILUTION =====================
    protocol.comment("")
    protocol.comment("=" * 50)
    protocol.comment("PHASE 3: ICP-MS 10x DILUTION")
    protocol.comment("20 uL supernatant + 180 uL 2% HNO3 = 200 uL total")
    protocol.comment("=" * 50)

    for col_idx in range(12):
        protocol.comment(f"Diluting column {col_idx + 1} of 12")
        p300_multi.pick_up_tip()

        # Transfer 20 uL supernatant to dilution plate
        p300_multi.aspirate(20, supernatant_plate.columns()[col_idx][0])
        p300_multi.dispense(20, dilution_plate.columns()[col_idx][0])

        # Add 180 uL of 2% HNO3
        p300_multi.aspirate(180, hno3_2pct)
        p300_multi.dispense(180, 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("")
    protocol.comment("=" * 50)
    protocol.comment("PROTOCOL COMPLETE")
    protocol.comment("=" * 50)
    protocol.comment("")
    protocol.comment("Total tips used:")
    protocol.comment("  Phase 1 (P1000): ~19 tips")
    protocol.comment("  Phase 2 (P300):  96 tips (12 cols x 8 ch)")
    protocol.comment("  Phase 3 (P300):  96 tips (12 cols x 8 ch)")
    protocol.comment("  TOTAL:           ~211 tips")
    protocol.comment("")
    protocol.comment("Next steps:")
    protocol.comment("1. Label ICP-MS dilution plate (slot 9):")
    protocol.comment("   Date, 'Iter2-Roasted-10xDil', 60C/120min")
    protocol.comment("2. Seal and store at 4 C until ICP-MS analysis")
    protocol.comment("3. Store supernatant plate (slot 8) as backup at 4 C")
    protocol.comment("4. For 0.50-1.00M H2SO4 wells with high REE signal,")
    protocol.comment("   prepare additional 100x dilution manually:")
    protocol.comment("   20 uL of 10x diluted + 180 uL 2% HNO3")
    protocol.comment("")
    protocol.comment("Plate layout reference:")
    protocol.comment("  Rows: A=R1(NaOH+400C) B=R2(NaOH+500C) C=R3(NaOH+600C) D=R4(NaOH+800C)")
    protocol.comment("        E=R5(Dir+400C)   F=R6(Dir+500C)  G=R7(Dir+600C)  H=R8(Dir+800C)")
    protocol.comment("  Cols: 1-2=0.05M  3-4=0.10M  5-6=0.25M  7-8=0.50M  9-10=1.00M  11-12=AcOH")
