import pygame
import workers, towers, ring
import random

boxcols=[(90,90,90),(128,0,0),(0,128,0),(0,0,128)]
class Grid():
	def __init__(self,x,y,spacing,w,h):
		self.x,self.y,self.spacing,self.w,self.h=x,y,spacing,w,h
		self.xe,self.ye=x+spacing*w,y+spacing*h
		self.highlight=None
		self.act=0
		self.menu=None
		self.mx,self.my=0,0
		self.towers=[] # I'll make this a property later
		self.workers=[]
	
	def cell(self,x,y):
		sx,sy,s,w,h=self.x,self.y,self.spacing,self.w,self.h
		if not (sx<x<sx+s*w and sy<y<sy+s*h): return None
		x-=sx; y-=sy
		return x//s,y//s
	
	def cell_to_real(self,x,y):
		if not (0<x<self.w and 0<y<self.h): return 0,0,0
		s=self.spacing
		return self.x+s*x,self.y+s*y,s
	
	def add_tower(self,x,y,*args):
		self.towers.append(towers.Tower(x,y,self,*args))
		self.calc_tower_mean()
	
	def add_worker(self,behaviour):
		self.workers.append(workers.Worker(int(random.random()*self.spacing*self.w+self.y),int(random.random()*self.spacing*self.h+self.y),behaviour,self))
	
	def calc_tower_mean(self):
		coords=list(map(lambda x:(x.x,x.y),self.towers))
		x=list(map(lambda x:x[0],coords))
		y=list(map(lambda x:x[1],coords))
		x=sum(x)/len(x); y=sum(y)/len(y)
		mx,my,offset=self.cell_to_real(x,y)
		self.mx,self.my=mx+offset/2,my+offset/2
	
	def draw(self,screen):
		x,y,s,w,h=self.x,self.y,self.spacing,self.w,self.h
		xe,ye=self.xe,self.ye
		pygame.draw.rect(screen,(90,90,45),(x,y,s*w,s*h))
		for i in range(h+1):
			pygame.draw.line(screen,'black',(x,i*s+y),(xe,i*s+y))
		for i in range(w+1):
			pygame.draw.line(screen,'black',(i*s+x,y),(i*s+x,ye))
		for tower in self.towers: tower.draw(screen)
		if self.highlight is not None:
			hx,hy,_=self.cell_to_real(*self.highlight)
			hl=pygame.Surface((s,s))
			hl.set_alpha(128)
			pygame.draw.rect(hl,boxcols[self.act],(0,0,s,s))
			screen.blit(hl,(hx,hy))
		for worker in self.workers: worker.draw(screen)
		if self.menu: self.menu.draw(screen)
	
	def update(self,events):
		x,y=pygame.mouse.get_pos()
		if not self.act:
			self.highlight=self.cell(x,y)
			if self.menu: self.menu=None
		elif self.menu is None and self.highlight is not None:
			self.menu=ring.Ring(x,y,self.spacing*2.5,[ring_wrap(self,'cannon',0),ring_wrap(self,'cannon',1),ring_wrap(self,'cannon',2),ring_wrap(self,'cannon',3)],self)
		for event in events:
			if event.type==pygame.KEYUP: self.act=0; continue
			if event.type!=pygame.KEYDOWN: continue
			if event.scancode not in [4,7,22]: continue
			self.act={4:1,22:2,7:3}[event.scancode]
		for tower in self.towers: tower.update(events)
		for worker in self.workers: worker.update(events)
		if self.menu: self.menu.update(events)

def ring_wrap(grid,ttype,elem):
	def ring_inner(x,y):
		grid.add_tower(*grid.cell(x,y),ttype,elem)
	return ring_inner