!pip install opentrons>=3.13.0  # Update to a newer version of Opentrons
from opentrons import simulate, types

metadata = {
    'protocolName': 'Kenya Opentrons Lab',
    'author': 'Machoka',
    'source': 'Kenya 2025',
    'apiLevel': '2.20'
}

# Robot deck setup constants
TIP_RACK_DECK_SLOT = 9
COLORS_DECK_SLOT = 6
AGAR_DECK_SLOT = 5

well_colors = {
    'A1': 'Black',
    'B1': 'Red',
    'C1': 'Green',
    'D1': 'White'
}

# Estimated number of drops per color
color_drops = {
    'Black': 45,  # Top stripe + part of "KENYA"
    'Red': 45,    # Middle stripe + part of "KENYA"
    'Green': 45,  # Bottom stripe + part of "KENYA"
    'White': 15    # Oval + spears
}

# Define the custom labware
custom_labware_def = {
    "metadata": {
        "displayName": "Kenya Agar Plate",
        "displayCategory": "other",
        "displayVolumeUnits": "mL"
    },
    "cornerOffsetFromSlot": {"x": 0, "y": 0, "z": 0},
    "parameters": {
        "isTiprack": False,
        "loadName": "kenya_agar_plate",
        "isMagneticModuleCompatible": False,
        "quirks": []
    },
    "ordering": [["A1"]],
    "wells": {
        "A1": {
            "depth": 20,
            "totalLiquidVolume": 20000,
            "shape": "circular",
            "diameter": 90,
            "x": 0, "y": 0, "z": 0
        }
    },
    "dimensions": {"xDimension": 127.75, "yDimension": 85.45, "zDimension": 20},
    "brand": {"brand": "Generic", "brandId": ["generic"]},
    "namespace": "custom_beta",
    "version": 1,
    "schemaVersion": 2,
    "category": "other"
}

def run(protocol):
    # Define the custom labware within the protocol context
    protocol.define_labware(custom_labware_def)

    # Load Tip Rack
    tips_20ul = protocol.load_labware('opentrons_96_tiprack_20ul',
                                       TIP_RACK_DECK_SLOT,
                                       'Opentrons 20uL Tips')

    # Load Pipette
    pipette_20ul = protocol.load_instrument("p20_single_gen2",
                                            "right",
                                            tip_racks=[tips_20ul])

    # Load Modules
    temperature_module = protocol.load_module('temperature module gen2',
                                              COLORS_DECK_SLOT)
    
    # Load Color Plate
    temperature_plate = temperature_module.load_labware(
        'opentrons_96_aluminumblock_generic_pcr_strip_200ul', 'Cold Plate')
    color_plate = temperature_plate

    # Load Agar Plate
    agar_plate = protocol.load_labware(
        custom_labware_def['parameters']['loadName'], AGAR_DECK_SLOT,
        'Agar Plate')  
    center_location = agar_plate['A1'].top()
    
    def location_of_color(color_string):
        for well, color in well_colors.items():
            if color.lower() == color_string.lower():
                return color_plate[well]
        raise ValueError(f"No well found with color {color_string}")
    
    def dispense_with_new_tip(pipette, volume, source, destination, repetitions):
        pipette.pick_up_tip()
        for _ in range(repetitions):
            pipette.aspirate(volume, source)
            pipette.dispense(volume, destination)
        pipette.drop_tip()
    
    # Dispense each color using a fresh tip for the estimated number of drops
    for color, drops in color_drops.items():
        source_well = location_of_color(color)
        dispense_with_new_tip(pipette_20ul, 2, source_well, center_location, drops)