# This program creates a grid, and randomly places a certain number of dead ants represented by
# blue squares and a number of live ants represented by red circles. Then repeats it until you close the window.

# The dead ants are placed in a matrix 'world' with id = 1. The live ants are maintained in a different matrix
# 'live_ant_tracker' with id = 2, as a matrix cell can only hold one value.

# NOTE: Please insert your logic into the place marked so. It's advisable that you keep the rest of the program
# intact, as it has mostly to do with wxPython, and disturbing those piece might mess up the graphics part.

# wxPython NOTE: Most of the code here is similar to '2Dcells.py' provided for Q2 in Lab 3.
# This program has two additional things: 'Timer' and 'PostEvent'. The flickering happens every half
# second, which is controlled by a 'Timer' whose event is bound to the 'Update' method. In that
# method, the cell is freshly painted. Once the painting is done, it explicitly generates a 'EVT_PAINT'
# event in order to invoke the 'OnPaint' method. In that method, a new timer is set that is stopped
# in the 'Update' method.

import wx
import numpy as np
import random

# Grid specifications
GRID_B = 640
GRID_H = 480

# Cell specifications
NUM_HORIZ_CELLS = 10
NUM_VERT_CELLS = 10
CELL_WIDTH = float(GRID_B) / NUM_HORIZ_CELLS
CELL_HEIGHT = float(GRID_H) / NUM_VERT_CELLS
CIRCLE_RADIUS = min(CELL_WIDTH, CELL_HEIGHT) / 2.0

# Ant specifications
NUM_DEAD_ANTS = 20
NUM_LIVE_ANTS = 5

def OnPaint(event):
    panel.dc = wx.PaintDC(panel)
    timer.Start(1000)

def Update(event):
    for i in range(0, NUM_VERT_CELLS):
        for j in range(0, NUM_HORIZ_CELLS):
            x_pos = j * CELL_WIDTH
            y_pos = i * CELL_HEIGHT

            # Create blue squares for dead ants
            if frame.world[i, j] == 1:
                col = "blue"
            else:                         
                col = "white"
            panel.dc.SetBrush(wx.Brush(col, wx.SOLID))
            panel.dc.DrawRectangle(x_pos, y_pos, CELL_WIDTH, CELL_HEIGHT)

            # Create red circles for live ants
            center_x = x_pos +  (CELL_WIDTH / 2.0)
            center_y = y_pos + (CELL_HEIGHT / 2.0)
            if frame.live_ant_tracker[i, j] == 2:
                col = 'red'
                panel.dc.SetBrush(wx.Brush(col, wx.SOLID))
                panel.dc.DrawCircle(center_x, center_y, CIRCLE_RADIUS)

    # ----------------- INSERT YOUR LOGIC HERE ----------------------------#

    Randomize_worlds()   #Comment this out

    # ----------------- INSERT YOUR LOGIC HERE ----------------------------#
      
    timer.Stop()
    wx.PostEvent(panel, wx.PaintEvent())

def Randomize_worlds():
    # First create an empty world, initialized with 0's
    frame.world = np.zeros([NUM_VERT_CELLS, NUM_HORIZ_CELLS],dtype=int)

    # Randomly place 'dead ants' with id = 1
    cell_indices = random.sample(range(NUM_HORIZ_CELLS * NUM_VERT_CELLS), NUM_DEAD_ANTS)
    row_indices = [idx % NUM_VERT_CELLS for idx in cell_indices]
    column_indices = [idx / NUM_VERT_CELLS for idx in cell_indices]
    frame.world[row_indices, column_indices] = 1

    # NOTE: Create a separate 'tracker matrix' for live ants. The reason for doing this is that if live ants step onto cells
    # containing dead ants (which they are allowed to), we want to be able to track that, as the 'world' contains information
    # only about dead ants.
    # As an alternative to maintaining a separate tracker, you can simply associate a position vector (x-y coordinates) 
    # with each ant as a property of it (in object-oriented language), and just a maintain a list of live ants.

    frame.live_ant_tracker = np.zeros([NUM_VERT_CELLS, NUM_HORIZ_CELLS],dtype=int)

    # Randomly place 'live ants' with id = 2
    cell_indices = random.sample(range(NUM_HORIZ_CELLS * NUM_VERT_CELLS), NUM_LIVE_ANTS)
    row_indices = [idx % NUM_VERT_CELLS for idx in cell_indices]
    column_indices = [idx / NUM_VERT_CELLS for idx in cell_indices]
    frame.live_ant_tracker[row_indices, column_indices] = 2

app = wx.App()
frame = wx.Frame(None, title="Act Clustering", size = (640, 500))
panel = wx.Panel(frame)
panel.SetBackgroundColour("black")

panel.Bind(wx.EVT_PAINT, OnPaint)
timer = wx.Timer(panel)
panel.Bind(wx.EVT_TIMER, Update, timer)

Randomize_worlds()
        
#frame.Center()
frame.Show(True)

# A mainloop is an endless cycle that catches up all events coming up to your application.
# It is an integral part of any windows GUI application.
app.MainLoop()