import Image, ImageTk # PIL
import time, sys
import Tkinter as tk
import program, debugger

"""
Team Sealab 2021
jack@performancedrivers.com
bob@syristatides.com

This module parses boards and launches the GUI interface
"""

CLEAR = 0x01
ROCK = 0x02
RED_HILL = 0x04
BLK_HILL = 0x08
Gred = (255,0,0)
Gblk = (0,200,200)
Gfood = (0,255,0)

IMG_SCALE = 1.75
MAX_TURNS = 100000
MAX_TURNS = 20000

def image_open(filename):
  img = Image.open(filename)
  if (IMG_SCALE > 1):
    (w, h) = img.size
    img = img.resize((w*IMG_SCALE, h*IMG_SCALE))
  return img

def fill(img, color, replace=(255,255,255)):
  (imgw, imgh) = img.size
  for (x) in range(imgw):
    for (y) in range(imgh):
      if (img.getpixel((x,y)) == replace):
        img.putpixel((x,y), color)

class Ant(object):
  color = None
  counter = 0 # id counter
  all = []
  def __init__(self, coord):
    self.dir = 0
    self.pos = coord
    self.state = 0
    self.resting = 0
    self.has_food = 0
    self.id = Ant.counter
    Ant.counter += 1
    self.dead = 0
    self.stepped = 0
    self.type = 0
    Ant.all.append(self) # DEBUG, make this a weakref
    return
  def __str__(self):
    return "%d: pos %s, state %d, resting %d, food %d, dead %d" % (self.id, repr(self.pos), self.state, self.resting, self.has_food, self.dead)

class BlkAnt(Ant):
  color = BLK_HILL
  image = image_open('ant_square.png')
  fill(image, Gblk)
  base_image = image
  fill_color = Gblk

class RedAnt(Ant):
  color = RED_HILL
  image = image_open('ant_square.png')
  fill(image, Gred)
  base_image = image
  fill_color = Gred

def dir_left(dir):
  bef = dir
  dir = (6 + dir  - 1) % 6
  #print "Left, bef/aft", (bef, dir)
  return dir
def dir_right(dir):
  bef = dir
  dir = (6 + dir  + 1) % 6
  #print "Right, bef/aft", (bef, dir)
  return dir

MAX_ANTS = 300
#MAX_ANTS = 1

class Board(object):
  def __init__(self, filename):
    self.steped = 0
    fob = open(filename)
    self.width = int(fob.readline().strip())
    self.height = int(fob.readline().strip())

    # create an empty map
    self.rows = map(lambda x:[], range(self.height)) # [[]] * height
    for (row) in self.rows:
      for (i) in range(self.width):
        row.append({'type':None,
                    'food':0,
                    'hill':None,
                    'ant':None,
                    'redmark':[0]*6,
                    'blkmark':[0]*6,
                   })
    
    for (y, line) in enumerate(fob):
      for (x, c) in enumerate(filter(lambda x: x!=' ', line.strip())):
        self[(x,y)]['type'] = CLEAR
        if (c == '#'):
          self[(x,y)]['type'] = ROCK
        elif (c == '.'): pass
        elif (c == '-'):
          self[(x,y)]['hill'] = BLK_HILL
          if (len(filter(lambda x:x.color == BLK_HILL, Ant.all)) < MAX_ANTS):
            self[(x,y)]['ant'] = BlkAnt((x,y))
        elif (c == '+' or c == 'X'):
          self[(x,y)]['hill'] = RED_HILL
          # testing
          if (1 and c == 'X'):
            self[(x,y)]['redmark'][3] = 1
          if (len(filter(lambda x:x.color == RED_HILL, Ant.all)) < MAX_ANTS):
            self[(x,y)]['ant'] = RedAnt((x,y))
        else:
          try:
            food = int(c)
          except ValueError:
            print "DOH %s" % (repr(c))
            continue
          self[(x,y)]['food'] = food
    return

  def scores(self):
    red = 0
    black = 0
    for (x) in range(self.width):
      for (y) in range(self.height):
        if (self[(x,y)]['food'] and self[(x,y)]['hill']):
          if (self[(x,y)]['hill'] == RED_HILL):
            red += self[(x,y)]['food']
          else:
            black += self[(x,y)]['food']
    return (red, black)

  def __getitem__(self, pos):
    (x, y) = pos
    return self.rows[y][x]
  def rel_pos(self, pos, dir):
    (x, y) = pos
    if (y % 2 == 1):
      if (dir == 0):
        goto = (x+1, y)
      elif (dir == 1):
        goto = (x+1, y+1)
      elif (dir == 2):
        goto = (x, y+1)
      elif (dir == 3):
        goto = (x-1, y)
      elif (dir == 4):
        goto = (x, y-1)
      elif (dir == 5):
        goto = (x+1, y-1)
    else:
      if (dir == 0):
        goto = (x+1, y)
      elif (dir == 1):
        goto = (x, y+1)
      elif (dir == 2):
        goto = (x-1, y+1)
      elif (dir == 3):
        goto = (x-1, y)
      elif (dir == 4):
        goto = (x-1, y-1)
      elif (dir == 5):
        goto = (x, y-1)
    return goto

G_marker = {}
def make_marker(img, info):
  global G_marker
  tup = tuple(info)
  if (tup not in G_marker):
    img = img.copy()
    for (v, color) in zip(tup, ((128,0,0), (255,0,0), (0,128,0), (0,255,0), (0,0,128), (0,0,255))):
      if (not v):
        fill(img, (255,255,255), color)
    G_marker[tup] = img
  else:
    img = G_marker[tup]
  return img

class PNGBoard(Board):
  def __init__(self, *args, **opts):
    Board.__init__(self, *args, **opts)
    self.square = image_open('square.png')
    self.blk_hill = image_open('square.png')
    fill(self.blk_hill, Gblk)
    self.red_hill = image_open('square.png')
    fill(self.red_hill, Gred)
    self.food = image_open('square.png')
    fill(self.food, Gfood)
    self.blk_ant = image_open('ant_square.png')
    fill(self.blk_ant, Gblk)
    self.marker = image_open('marker.png')
    (sq_w, sq_h) = self.square.size
    big_w = self.width * (sq_w -1) + sq_w-2
    big_h = self.height * (sq_h -1) + 1
    self.img = Image.new('RGB', (big_w, big_h))
    self.redmark = Image.new('RGB', (big_w, big_h))
    self.blkmark = Image.new('RGB', (big_w, big_h))
    self.remake_img()
    return
  
  def remake_img(self):
    for (y) in range(self.height):
      for (x) in range(self.width):
        self.set_square((x,y), self[(x,y)], 0)
    return

  def set_square(self, coord, info, save=1):
    (sq_w, sq_h) = self.square.size
    (x, y) = coord
    if (y > 0):
      curry = y * (sq_h-1)
    else:
      curry = 0
    if (x > 0):
      currx = x * (sq_h-1)
    else:
      currx = 0
    if (y % 2 == 1):
      currx += sq_w / 2

    t = info['type']
    h = info['hill']
    if (info['ant']):
      ant = info['ant']
      square = ant.image
    elif (info['food']):
      square = self.food
    elif (h == BLK_HILL):
      square = self.blk_hill
    elif (h == RED_HILL):
      square = self.red_hill
    elif (t == ROCK):
      return
    else:
      square = self.square

    self.img.paste(square, (currx, curry))

    mark = make_marker(self.marker, info['blkmark'])
    self.blkmark.paste(mark, (currx, curry))

    mark = make_marker(self.marker, info['redmark'])
    self.redmark.paste(mark, (currx, curry))
    return

class TKBoard(PNGBoard):
  def __init__(self, root, *args, **opts):
    print (self, root, args, opts)
    PNGBoard.__init__(self, *args, **opts)
    self.__tk = root
    self.can = tk.Canvas(root)    
    self.can.pack(fill=tk.BOTH)
    (img_w, img_h) = self.img.size
    self.can.config(width=img_w, height=img_h)
    self.tkdelme = None
    self.tkimg = None
    self.which = 'main' # main, black, red
    self.but1 = tk.Button(root,text='Main', command=lambda :self.set_which('main'))
    self.but1.pack(side=tk.LEFT)
    self.but2 = tk.Button(root,text='Black Mark', command=lambda :self.set_which('black'))
    self.but2.pack(side=tk.LEFT)
    self.but3 = tk.Button(root,text='Red Mark', command=lambda :self.set_which('red'))
    self.but3.pack(side=tk.LEFT)
    self.refresh()
    return

  def cleanup(self):
    self.can.delete(self.tkdelme)
    self.but1.destroy()
    self.but2.destroy()
    self.but3.destroy()
    self.can.destroy()

  def set_which(self, which):
    self.which = which
    self.refresh(0)
    return

  def refresh(self, reschedule=1):
    self.remake_img()
    if (not self.tkdelme is None):
      self.tkimg = None
      self.can.delete(self.tkdelme)
      self.tkdelme = None

    if (self.which == 'main'):
      img = self.img
    elif (self.which == 'red'):
      img = self.redmark
    elif (self.which == 'black'):
      img = self.blkmark
    else:
      raise Exception("Wrong thingy")
    self.tkimg = ImageTk.PhotoImage(img)
    self.tkdelme = self.can.create_image(2, 2, image=self.tkimg, anchor=tk.NW)
    #if (reschedule):
    #  self.__tk.after(500, self.refresh) # reschedule ourselves to happen again very soon    
    return

class MDropDown(tk.Frame):
  def __init__(self, root, labels, default = None):
    tk.Frame.__init__(self, root)
    self.__val = tk.StringVar()
    if (default and default in labels):
      self.__val.set(default)
    else:
      self.__val.set(labels[0])
    self.__opt = tk.OptionMenu(self, self.__val, *labels)
    self.repack()

  def repack(self):
    self.__opt.pack(side=tk.TOP, fill=tk.X)
    self.pack(side=tk.TOP)
    return

  def value(self, set=None):
    if (set is not None):
      self.__val.set(set)
    return self.__val.get()
  

class BigTk(object):
  boards = ['sample0.world','sample1.world','sample2.world','sample3.world','sample4.world','sample5.world','sample6.world','sample7.world','sample8.world','sample9.world','tiny.world',
            'gathertest.world',]
  progs = ['pher.ant', 'pher2.ant', 'pher3.ant', 'pher_erase.ant', 'pher_erase2.ant', 'test.ant', 'figure.ant']
  def __init__(self, *ignored):
    self.tk = tk.Tk()
    self.tk.title('Ant Simulator')
    self.curr_board_name = self.boards[0]
    print "CBN", self.curr_board_name
    self.cmds = tk.Frame(self.tk)
    self.cmds.pack(side=tk.LEFT)
    self.dropdown = MDropDown(self.cmds, self.boards, 'sample3.world')
    self.redprog = MDropDown(self.cmds, self.progs, 'pher2.ant')
    self.blkprog = MDropDown(self.cmds, self.progs, 'pher2.ant')
    self.revbut = tk.Button(self.cmds,text='Reverse', command=self.reverse)
    self.revbut.pack(side=tk.TOP)
    self.resbut = tk.Button(self.cmds,text='Reset', command=self.reset)
    self.resbut.pack(side=tk.TOP)
    self.step1 = tk.Button(self.cmds,text='Step', command=self.step)
    self.step1.pack(side=tk.TOP)
    self.step10 = tk.Button(self.cmds,text='Step10', command=self.step10)
    self.step10.pack(side=tk.TOP)
    self.step100 = tk.Button(self.cmds,text='Step100', command=self.step100)
    self.step100.pack(side=tk.TOP)
    self.step1000 = tk.Button(self.cmds,text='Step1k', command=self.step1000)
    self.step1000.pack(side=tk.TOP)
    self.runbut = tk.Button(self.cmds,text='Run', command=self.run)
    self.runbut.pack(side=tk.TOP)
    self.quitbut = tk.Button(self.cmds,text='Quit', command=self.tk.quit)
    self.quitbut.pack(side=tk.TOP)
    self.curr_board = TKBoard(self.tk, self.curr_board_name)
    self.rprogram = None
    self.bprogram = None
    self.curr_red = self.redprog.value()
    self.curr_blk = self.blkprog.value()
    self.debug = None
    self.check_dropdown()
    return

  def reset(self):
    self.curr_red = 'bogus'
    return

  def reverse(self):
    red = self.redprog.value()
    blk = self.blkprog.value()
    self.blkprog.value(red)
    self.redprog.value(blk)
    self.curr_red = 'bogus'
    return

  def check_dropdown(self):
    newv = self.dropdown.value()
    changed = 0
    if (newv != self.curr_board_name):
      self.curr_board_name = newv
      changed = 1
    elif (self.redprog.value() != self.curr_red):
      self.curr_red = self.redprog.value()
      changed = 1
    elif (self.blkprog.value() != self.curr_blk):
      self.curr_blk = self.blkprog.value()
      changed = 1
    if (changed):
      self.curr_board.cleanup()
      Ant.all = [] # RESET
      self.curr_board = TKBoard(self.tk, newv)
      self.rprogram = program.SuperProgram(self.curr_board, self.curr_red, tkroot=self.tk)
      if (self.debug):
        self.debug.destroy()
      self.debug = debugger.Debugger(self.tk, self.rprogram.intermed, self.curr_red)
      
      self.bprogram = program.SuperProgram(self.curr_board, self.curr_blk, tkroot=self.tk)
    self.tk.after(500, self.check_dropdown)
    return

  def step10(self):
    return self.step(10)
  def step100(self):
    return self.step(100)
  def step1000(self):
    return self.step(1000)
  def step(self, cnt=1):
    some = 0
    t = time.time()
    allants = Ant.all
    max_turns = MAX_TURNS
    red_hill = RED_HILL
    while (cnt and self.curr_board.steped < max_turns):
      cnt -= 1
      self.curr_board.steped += 1
      for (ant) in allants:
        if (ant.color == red_hill):
          self.rprogram.execute(ant)
        else:
          self.bprogram.execute(ant)
        if (len(allants) == 1):
          self.debug.display_state(ant.state)
      if (some % 5000 == 0):
        sys.stderr.write("STEPS %d %4.2f\n" % (self.curr_board.steped, time.time()-t))
      #  self.curr_board.refresh()
      some += 1
    self.curr_board.refresh()
    (red, blk) = self.curr_board.scores()
    print "RED %d %s" % (red, repr(self.curr_red))
    print "BLACK %d %s" % (blk, repr(self.curr_blk))
    return
  def run(self):
    t = time.time()
    self.step(-1)
    print "TIME", time.time() - t
    return

def test():
  #b = BigTk('test.ant')
  #b = BigTk('sample.ant')
  #b = BigTk('pher2.ant', 'sample.ant')  
  #b = BigTk('sample.ant', 'pher2.ant')
  # pre psycho
  #TIME 32.9414381981
  #RED/BLACK (55, 26)

  #b = BigTk('pher_erase.ant', 'pher_erase2.ant')
  b = BigTk()
  b.tk.mainloop()

if (__name__ == '__main__'):
  import profile
  #profile.run('test()')
  #import psyco
  #psyco.full()
  test()
