from opentrons import types

metadata = {
    'protocolName': 'HTGAA Opentrons Lab',
    'author': 'HTGAA',
    'source': 'HTGAA 2022',
    'apiLevel': '2.20'
}

##############################################################################
###   Robot deck setup constants - don't change these
##############################################################################

TIP_RACK_DECK_SLOT = 9
COLORS_DECK_SLOT = 6
AGAR_DECK_SLOT = 5
PIPETTE_STARTING_TIP_WELL = 'A1'

well_colors = {
    'A1' : 'Red',
    'B1' : 'Yellow',
    'C1' : 'Green',
    'D1' : 'Cyan',
    'E1' : 'Blue'       # if in a 24-well plate, this needs to be moved to e.g. D2
}

def run(protocol):
  ##############################################################################
  ###   Load labware, modules and pipettes
  ##############################################################################

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

  # Pipettes
  pipette_20ul = protocol.load_instrument("p20_single_gen2", "right", [tips_20ul])

  # Modules
  temperature_module = protocol.load_module('temperature module gen2', COLORS_DECK_SLOT)

  # Temperature Module Plate
  temperature_plate = temperature_module.load_labware('opentrons_96_aluminumblock_generic_pcr_strip_200ul',
                                                      'Cold Plate')
  # Choose where to take the colors from
  color_plate = temperature_plate

  # Agar Plate
  agar_plate = protocol.load_labware('htgaa_agar_plate', AGAR_DECK_SLOT, 'Agar Plate')  ## TA MUST CALIBRATE EACH PLATE!
  # Get the top-center of the plate, make sure the plate was calibrated before running this
  center_location = agar_plate['A1'].top()

  pipette_20ul.starting_tip = tips_20ul.well(PIPETTE_STARTING_TIP_WELL)

  ##############################################################################
  ###   Patterning
  ##############################################################################

  ###
  ### Helper functions for this lab
  ###

  # pass this e.g. 'Red' and get back a Location which can be passed to aspirate()
  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}")

  # For this lab, instead of calling pipette.dispense(1, loc) use this: dispense_and_jog(pipette, 1, loc)
  def dispense_and_jog(pipette, volume, location):
      """
      Dispense and then move up 5mm and back down to shake all dispensed fluid off the tip;
      this also ensures it's not moving laterally before the dispense is done.
      """
      assert(isinstance(volume, (int, float)))
      pipette.dispense(volume, location)
      currLoc = pipette._get_last_location_by_api_version()
      pipette.move_to(currLoc.move(types.Point(z=5)))
      pipette.move_to(currLoc)

  ###
  ### YOUR CODE HERE to create your design
  ###

  # H
  cursor = center_location.move(types.Point(x=-35.5, y = 12))

  pipette_20ul.pick_up_tip()

  pipette_20ul.aspirate(20, location_of_color('Red'))
  for i in range(12):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(x=2)))
    if i == 9:
      pipette_20ul.aspirate(10, location_of_color('Red'))
    cursor = cursor.move(types.Point(y =-2))

  cursor = cursor.move(types.Point(x=4, y=14))

  for i in range(3):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(y=-2)))
    cursor = cursor.move(types.Point(x =2))

  cursor = cursor.move(types.Point(y=10))

  pipette_20ul.aspirate(20, location_of_color('Red'))
  for i in range(12):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(x=2)))
    cursor = cursor.move(types.Point(y =-2))

    if i == 9:
      pipette_20ul.aspirate(4, location_of_color('Red'))

  cursor = cursor.move(types.Point(x=4, y=24))

  pipette_20ul.drop_tip()
  pipette_20ul.pick_up_tip()

  # T
  pipette_20ul.aspirate(16, location_of_color('Yellow'))
  for i in range(8):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(y=-2)))
    cursor = cursor.move(types.Point(x =2))

  cursor = cursor.move(types.Point(x=-10, y=-4))

  pipette_20ul.aspirate(20, location_of_color('Yellow'))
  for i in range(10):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(x=2)))
    cursor = cursor.move(types.Point(y =-2))

  pipette_20ul.drop_tip()
  pipette_20ul.pick_up_tip()

  # G
  cursor = cursor.move(types.Point(x=10, y=2))

  pipette_20ul.aspirate(20, location_of_color('Green'))
  for i in range(12):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(x=2)))
    cursor = cursor.move(types.Point(y =2))

    if i == 9:
      pipette_20ul.aspirate(14, location_of_color('Green'))

  cursor = cursor.move(types.Point(x=4, y=-2))

  for i in range(5):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(y=-2)))
    cursor = cursor.move(types.Point(x =2))

  cursor = cursor.move(types.Point(x=-4, y=-12))

  pipette_20ul.aspirate(18, location_of_color('Green'))

  for i in range(6):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(x=2)))
    cursor = cursor.move(types.Point(y=-2))

  cursor = cursor.move(types.Point(x=-2, y=4))

  for i in range(3):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(y=-2)))
    cursor = cursor.move(types.Point(x =-2))

  # A

  pipette_20ul.drop_tip()
  pipette_20ul.pick_up_tip()

  cursor = cursor.move(types.Point(x=12, y=-2))

  pipette_20ul.aspirate(20, location_of_color('Cyan'))
  for i in range(12):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(x=2)))
    cursor = cursor.move(types.Point(y =2, x=0.4))

    if i == 9:
      pipette_20ul.aspirate(4, location_of_color('Cyan'))

  cursor = cursor.move(types.Point(x=2, y=-2))

  pipette_20ul.aspirate(20, location_of_color('Cyan'))
  for i in range(12):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(x=2)))
    cursor = cursor.move(types.Point(y =-2, x=0.4))

    if i == 9:
      pipette_20ul.aspirate(4, location_of_color('Cyan'))

  cursor = cursor.move(types.Point(x=4, y=2))

  # A

  pipette_20ul.drop_tip()
  pipette_20ul.pick_up_tip()

  pipette_20ul.aspirate(20, location_of_color('Blue'))
  for i in range(12):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(x=2)))
    cursor = cursor.move(types.Point(y =2, x=0.4))

    if i == 9:
      pipette_20ul.aspirate(4, location_of_color('Blue'))

  cursor = cursor.move(types.Point(x=2, y=-2))

  pipette_20ul.aspirate(20, location_of_color('Blue'))
  for i in range(12):
    dispense_and_jog(pipette_20ul, 1, cursor)
    dispense_and_jog(pipette_20ul, 1, cursor.move(types.Point(x=2)))
    cursor = cursor.move(types.Point(y =-2, x=0.4))

    if i == 9:
      pipette_20ul.aspirate(4, location_of_color('Blue'))

  pipette_20ul.drop_tip()