from opentrons import protocol_api

# Define metadata for the protocol
metadata = {
    "protocolName": "Liquid Dispensing for 96-Well Plate",
    "description": """This protocol dispenses varying amounts of NaOH and Produced Water,
                   plus a fixed amount of DI Water, into a 96-well plate according
                   to a specified layout. It is designed for an Opentrons Flex.
                   Total volume per well is 360 µL.""",
    "author": "Opentrons AI",
    "apiLevel": "2.16"
}

# Define requirements for the protocol
requirements = {"robotType": "Flex"}

def run(protocol: protocol_api.ProtocolContext):
    """
    Main function of the protocol.
    """
    # --- Labware and Instrument Loading ---
    protocol.comment("Loading labware and instruments...")

    # Load a 12-well reservoir for source liquids on deck slot D2
    reservoir = protocol.load_labware(
        load_name="nest_12_reservoir_15ml",
        location="D2"
    )

    # Load the 96-well plate on deck slot D3
    plate = protocol.load_labware(
        load_name="axygen_96_wellplate_500ul",
        location="D3"
    )

    # Load a tip rack for the pipette on deck slot D1
    # For improved efficiency with volumes > 200 µL,
    # consider using a 1000 µL tip rack ("opentrons_flex_96_tiprack_1000ul").
    tiprack = protocol.load_labware(
        load_name="opentrons_flex_96_tiprack_1000ul",
        location="D1"
    )

    # Load the Flex 1-Channel 1000 uL pipette on the left mount
    pipette = protocol.load_instrument(
        instrument_name="flex_1channel_1000",
        mount="left",
        tip_racks=[tiprack]
    )

    # --- Source Liquid Definitions ---
    # Assign liquids to wells in the reservoir
    naoh_source = reservoir["A1"]
    pw_source = reservoir["A2"]
    di_water_source = reservoir["A3"]
# Load trash bin
    trash_bin = protocol.load_trash_bin("A3")

    # --- Volume Definitions ---
    # Volumes of NaOH and PW for each row (A-H)
    naoh_volumes = [30, 60, 90, 120, 150, 180, 210, 0]
    pw_volumes = [300, 270, 240, 210, 180, 150, 120, 330]
    di_water_volume = 30

    # --- Dispensing Procedure ---

    # 1. Distribute DI Water
    # Add 30 µL of DI Water to all 96 wells
    protocol.comment(f"Adding {di_water_volume} µL of DI Water to all wells.")
    pipette.distribute(di_water_volume, di_water_source, plate.wells())

    # 2. Distribute NaOH and Produced Water (PW) row by row
    for i in range(8):
        target_row = plate.rows()[i]
        row_name = chr(ord('A') + i)

        # Get the specific volumes for the current row
        naoh_vol = naoh_volumes[i]
        pw_vol = pw_volumes[i]

        # Dispense NaOH for the current row, if volume is greater than 0
        if naoh_vol > 0:
            protocol.comment(f"Adding {naoh_vol} µL of NaOH to row {row_name}.")
            pipette.distribute(naoh_vol, naoh_source, target_row)

        # Dispense Produced Water (PW) for the current row
        protocol.comment(f"Adding {pw_vol} µL of PW to row {row_name}.")
        pipette.distribute(pw_vol, pw_source, target_row)

    protocol.comment("Protocol complete.")
