import supply
import random
import matrixbot
import numpy as np
from operator import itemgetter

SOLD_BENEFIT = 2

def run_generation(pop, i):
    print("Generation %d" % i)
    ranking = []
    for x in pop:
        score = evaluate(x)
        ranking.append((score, x))
    ranking.sort(key=itemgetter(0))
    print("Best of generation: ", ranking[0])
    newpop = [x[1] for x in ranking[:len(ranking)//2]]
    newpop = breed(newpop, len(pop), i)
    return newpop

def evaluate(mem):
    g = supply.variableRequestGenerator(supply.TURNS, 8, 14)
    #chain = [StupidAI(12), PlayerChain("Manufacturer", 12), Base()]
    chain = [matrixbot.MatrixBot(mem[:25], mem[25:], 12), supply.Base()]
    total_sold = 0
    history = []
    for topreq in g:
        requests = [x.request() for x in chain]
        requests = requests[:-1]
        requests.insert(0, topreq)
        got = 0
        for i, c in reversed(list(enumerate(chain))):
            got = c.reconcile(requests[i], got)
        #print("Sold ", got, " units.")
        total_sold += got
        history.append([x.cost() for x in chain])
    total_cost = [sum([x[0] * supply.SUPPLY_COST + x[1] * supply.BACKLOG_COST for x in y]) for y in history]
    cost = supply.integrate(total_cost)[-1] - (total_sold * SOLD_BENEFIT)
    return cost


def mate(a, b):
    out = []
    for x in range(len(a)):
        if random.random() < 0.5:
            out.append(a[x])
        else:
            out.append(b[x])
    return np.array(out)

def mutate(a):
    for i in range(5):
        x = random.randint(0, len(a) - 1)
        amt = random.uniform(-0.5, 0.5)
        a[x] = a[x] + amt

WEIGHT = 0.9
LATE_WEIGHT = 0.2
GENS = 500

def breed(pop, n, i):
    global WEIGHT
    mut_chance = (WEIGHT - LATE_WEIGHT) * (GENS - i / GENS) + LATE_WEIGHT
    total = n - len(pop)
    newbies = []
    while total != 0:
        a = random.choice(pop)
        b = random.choice(pop)
        c = mate(a, b)
        if random.random() < mut_chance:
            mutate(c)
        newbies.append(c)
        total -= 1
    c = pop[:]
    c.extend(newbies)
    return c

def main():
    starting_pop = []
    for i in range(300):
        x = np.random.uniform(low=-2, high=2, size=30)
        starting_pop.append(x)
    pop = starting_pop
    for gen in range(GENS):
        pop = run_generation(pop, gen+1)

if __name__ == "__main__":
    main()