forestGen2.py

import numpy as np
import matplotlib as ml
import matplotlib.pyplot as plt
from matplotlib import mpl,pyplot
import numpy as np
import thread

import sys,os,time,copy,math
from random import randrange

class Map:
    def __init__(self,size):
        #Generates map.
        #Generates a forest based on the attributes#
        #Lumberjack - 10% #
        #Trees - 
        #Bears - 
        self.map =  [["." for x in xrange(size)] for x in xrange(size)]
        self.lumberjacks = []
        self.trees = []
        self.bears = []
        self.tickCount = 0
        self.yearlyLumber = 0
        self.monthlyLog = []
        self.yearlyLog = []
        self.monthlySaplings = 0
        #Spawns lumberjacks - 10%#
        for i in range(0,int(round(((size*size)*0.1)))):
            x,y = findLocation(self.map)
            lumberjack = Lumberjack(x,y)
            self.lumberjacks.append(lumberjack)
            self.calculateMap(size)
        #Spawns trees - 50% #
        for i in range(0,int(round(((size*size)*0.5)))):
            x,y = findLocation(self.map)
            tree = Tree("Tree",x,y)
            self.trees.append(tree)
            self.calculateMap(size)
        #Spawns bears - 2% #
        for i in range(0,int(round(((size*size)*0.02)))):
            x,y = findLocation(self.map)           
            bear = Bear(x,y)
            self.bears.append(bear)
            self.calculateMap(size)
    def displayMap(self):
        #Displays the map as 2D ASCII#
        mapDisplay = ""
        for x in self.map:
            for y in x:
                mapDisplay = mapDisplay + str(y)
            mapDisplay = mapDisplay + "\n"
        print(mapDisplay)
    def visualMap(self):
        # make a color map of fixed colors
        cmap = mpl.colors.ListedColormap(['white','#75FF81','#05F719','#03910F','#FF0D0D','#5C400C'])
        bounds=[0,1,2,3,4,5,6]
        norm = mpl.colors.BoundaryNorm(bounds, cmap.N)

        # tell imshow about color map so that only set colors are used
        img = pyplot.imshow(self.visualise,interpolation='nearest',cmap = cmap,norm=norm)

        # make a color bar
        cbar = pyplot.colorbar(img,cmap=cmap,norm=norm,boundaries=bounds,ticks=[0.5,1.5,2.5,3.5,4.5,5.5,6.5])
        cbar.ax.set_yticklabels(['Nothing', 'Sapling','Tree','Elder','Lumberjack','Bear'])
        #pyplot.lcolorbar(("Nothing","Sapling","Tree","Elder","Lumberjack","Bear"))
        plt.ion()
        pyplot.show()
        while(True):
                plt.title("Year "+str(int(self.tickCount / 12))+ ", Month "+ str((self.tickCount % 12)+1))
                cmap = mpl.colors.ListedColormap(['white','#75FF81','#05F719','#03910F','#FF0D0D','#5C400C'])
                bounds=[0,1,2,3,4,5,6]
                norm = mpl.colors.BoundaryNorm(bounds, cmap.N)
                
                # tell imshow about color map so that only set colors are used
                img = pyplot.imshow(self.visualise,interpolation='nearest',cmap = cmap,norm=norm)
                if(self.drawable == 1):
                    plt.pause(0.5)
                #plt.draw()

    def calculateMap(self,size):
        self.map =  [["." for x in xrange(size)] for x in xrange(size)]
        self.visualise = [[0 for x in xrange(size)] for x in xrange(size)]
        self.drawable = 0
        for tree in self.trees:
            x,y = tree.getLocation()
            self.map[y][x] = tree          
            self.visualise[y][x] = tree.getDisplay()
        for lumberjack in self.lumberjacks:
            x,y = lumberjack.getLocation()
            self.map[y][x] = lumberjack
            self.visualise[y][x] = 4
        for bear in self.bears:
            x,y = bear.getLocation()
            self.map[y][x] = bear
            self.visualise[y][x] = 5
        self.drawable = 1
        

    
    def tick(self):
        self.tickCount += 1
        self.monthlyLog = []
        #Tree tick
        for tree in self.trees:
            #Drop sapling
            drop = tree.dropChance(self.map)
            if(drop != None):
                tree = Tree("Sapling",drop[0],drop[1])
                self.trees.append(tree)
                self.calculateMap(len(self.map[0]))
                self.monthlySaplings += 1
            #Age up
            type = tree.monthUp()
        self.monthlyLog.append("["+str(self.monthlySaplings) + "] new saplings.")
        self.monthlySaplings = 0
        lumber = 0
        #Lumberjack tick
        for lumberjack in self.lumberjacks:
            noOfLumber = lumberjack.move(self.map)
            if(noOfLumber > 0):
                lumber = lumber + noOfLumber
                self.yearlyLumber += noOfLumber
                location = lumberjack.getLocation()
                for i, tree in enumerate(self.trees):
                    if(tree.getLocation() == location):
                        del self.trees[i]
                        self.calculateMap(len(self.map[0]))
        #Bear tick
        for bear in self.bears:
            if(bear.move(self.map) == 1):
                location = bear.getLocation()
                for i, lumberjack in enumerate(self.lumberjacks):
                    if(lumberjack.getLocation() == location):
                        del self.lumberjacks[i]
                        self.calculateMap(len(self.map[0]))
        # Each year #
        if(self.tickCount % 12 == 0):
            self.yearlyLog = []
            self.yearlyLog.append("["+str(self.yearlyLumber)+"] pieces of lumber harvested by Lumberjacks.")
            if(self.yearlyLumber > len(self.lumberjacks)):
                amountToHire = int((self.yearlyLumber - len(self.lumberjacks)) / 10)
                for i in range(0,amountToHire):
                    x,y = findLocation(self.map)
                    lumberjack = Lumberjack(x,y)
                    self.lumberjacks.append(lumberjack)
                    #self.yearlyLog.append("Added jack at "+str(x)+","+str(y))
                    self.calculateMap(len(self.map[0]))
                self.yearlyLog.append("["+str(amountToHire)+"] new lumberjack(s)")
            elif(self.yearlyLumber < len(self.lumberjacks)):
                if(len(self.lumberjacks) > 1):
                    amountToFire = int(math.ceil((len(self.lumberjacks) - self.yearlyLumber) / 10))
                    for i in range(0,amountToFire):
                        del self.lumberjacks[random(len(self.lumberjacks))]
                        self.yearlyLog.append("A lumberjack has been fired")
            self.yearlyLog.append("["+str(len(self.lumberjacks))+"] Lumberjacks.")
            #Bear year#
            mawCount = 0
            for bear in self.bears:
                mawCount += bear.getMaws()
                bear.newYear()
            self.yearlyLog.append("["+str(mawCount)+"] maws.")
            if(mawCount > 0):
                del self.bears[random(len(self.bears))]
                self.yearlyLog.append("[1] bear has been captured and taken to the Zoo.")
                if(len(self.lumberjacks) == 0):
                    x,y = findLocation(self.map)
                    lumberjack = Lumberjack(x,y)
                    self.lumberjacks.append(lumberjack)
            elif(mawCount == 0):
                x,y = findLocation(self.map)
                bear = Bear(x,y)
                self.bears.append(bear)
                self.yearlyLog.append("[1] bear has roamed into the area.")
            ## Log information. ##            
        self.monthlyLog.append("["+str(lumber)+"] pieces of lumber harvested by Lumberjacks.")
          
    def log(self):        
        mawCount = 0
        print("Year "+str(int(self.tickCount / 12))+ ", Month "+ str((self.tickCount % 12)+1))
        print("========Monthly Log========")
        for logItem in self.monthlyLog:
            print(logItem)
        print("\n========Yearly log========")
        for logItem in self.yearlyLog:
            print(logItem)
        if(self.tickCount % 12 == 0):
            self.yearlyLumber = 0
        

class Tree:
    def __init__(self,treeType,x,y):
        self.type = treeType
        self.x = x
        self.y = y
        if(treeType == "Sapling"):
            self.age = 0
        elif(treeType == "Tree"):
            self.age = 12
        elif(treeType == "Elder"):
            self.age = 120
            
    def __str__(self):
        if(self.type == "Sapling"):
            return "1"
        elif(self.type == "Tree"):
            return "2"
        else:
            return "3"
    def getDisplay(self):
        if(self.type == "Sapling"):
            return 1
        elif(self.type == "Tree"):
            return 2
        elif(self.type == "Elder"):
            return 3

    def monthUp(self):
        self.age = self.age + 1
        if(self.age >= 120):
            self.type = "Elder"
        elif(self.age >= 12):
            self.type = "Tree"
    
    def drop(self,map,maxSize):
        validDrop = []
        for x in range(-1,2):
            for y in range(-1,2):
                checkX = self.x + x
                checkY = self.y + y
                if(checkX != -1 and checkX != maxSize):
                    if(checkY != -1 and checkY != maxSize):
                        if(map[checkY][checkX] == "."):
                            validDrop.append((checkX,checkY))
        if(len(validDrop) == 0):
            return None
        else:
            index = random(len(validDrop))
            return(validDrop[index])
    
    def getLocation(self):
        return self.x,self.y

    def dropChance(self,map):
        maxSize = len(map[0])
        randomNo = random(100)
        if(self.type == "Tree"):
            if(randomNo < 10):
                return self.drop(map,maxSize)
        if(self.type == "Elder"):            
            if(randomNo < 20):
                return self.drop(map,maxSize)

class Lumberjack:
    def __init__(self,x,y):
        self.x = x
        self.y = y
        self.woodCount = 0
        
    def __str__(self):
        return "4"
    
    def move(self,map):
        maxSize = len(map[0])
        validMove = []
        for i in range(0,3):
            for x in range(-1,2):
                for y in range(-1,2):
                    checkX = self.x + x
                    checkY = self.y + y
                    if(checkX != -1 and checkX != maxSize):
                        if(checkY != -1 and checkY != maxSize):
                            if(str(map[checkY][checkX]) != "S"):
                                validMove.append((checkX,checkY))
            index = random(len(validMove))
            self.x = validMove[index][1]
            self.y = validMove[index][0]
            if(str(map[self.y][self.x]) == "2"):
                self.woodCount = self.woodCount + 1
                return 1
            elif(str(map[self.y][self.x]) == "3"):
                self.woodCount = self.woodCount + 2
                return 2
            
            
    def getLocation(self):
        return self.x,self.y

class Bear:
    def __init__(self,x,y):
        self.x = x
        self.y = y
        self.mawCount = 0
    def __str__(self):
        return "5"
    def newYear(self):
        self.mawCount = 0
    def getMaws(self):
        return self.mawCount
    def move(self,map):
        maxSize = len(map[0])
        validMove = []
        for i in range(0,5):
            for x in range(-1,2):
                for y in range(-1,2):
                    checkX = self.x + x
                    checkY = self.y + y
                    if(checkX != -1 and checkX != maxSize):
                        if(checkY != -1 and checkY != maxSize):
                            validMove.append((checkX,checkY))
            index = random(len(validMove))
            self.x = validMove[index][1]
            self.y = validMove[index][0]
            if(str(map[self.y][self.x]) == "4"):
                if(random(100) < 80):
                    self.mawCount += 1
                    return 1
                else:
                    return 0
    def getLocation(self):
        return self.x,self.y
        
def findLocation(map):
    x = random(len(map[0]))
    y = random(len(map[0]))
    while(str(map[y][x]) != "."):
        x = random(len(map[0]))
        y = random(len(map[0]))
    return x,y

def random(size):
    return randrange(0,size)

def tick():
    print("Tick")
try:
    if __name__ == "__main__":
        map = Map(int(sys.argv[1]))        
        thread.start_new_thread(map.visualMap,())
        while(True):
            map.tick()
            os.system("cls")
            #map.displayMap()            
            map.calculateMap(int(sys.argv[1]))
            map.log()
            time.sleep(1)
except KeyboardInterrupt:
    sys.exit("\nFinished simulation")