import random
import time


class Unit:
    def __init__(self, name, e_cost, e_prod, bt, bp):
        self.name = name
        self.e_cost = e_cost
        self.e_prod = e_prod
        self.bt = bt
        self.bp = bp

    def bp_per_e(self):
        return self.bp / self.e_cost

    def e_per_e(self):
        return self.e_prod / self.e_cost

    def bp_per_bt(self):
        return self.bp / self.bt

    def e_per_bt(self):
        return self.e_prod / self.bt

    def is_fac(self):
        return 'f' in self.name

    def is_engy(self):
        return 'e' in self.name


acu_t1 = Unit('a1', 5000000, 20, 6000000, 10)
acu_t2 = Unit('a2', 21000, 0, 1000, 42-acu_t1.bp)                     # consumes an acu_t1
acu_t3 = Unit('a3', 50000, 0, 8400, 100-acu_t2.bp-acu_t1.bp)          # consumes an acu_t2

fac_navy_t1 = Unit('f1', 1500, 0, 300, 20)
fac_navy_t2_hq = Unit('f2h', 10000, 0, 3600, 90-fac_navy_t1.bp)        # consumes a fac_navy_t1
fac_navy_t3_hq = Unit('f3h', 45000, 0, 11250, 150-fac_navy_t2_hq.bp)   # consumes a fac_navy_t2_hq
fac_navy_t2_sl = Unit('f2s', 5500, 0, 3300, 90)
fac_navy_t3_sl = Unit('f3s', 11000, 0, 3500, 150)

engy_t1 = Unit('e1', 260, 0, 260, 5)
engy_t2 = Unit('e2', 650, 0, 650, 12.5)
engy_t3 = Unit('e3', 1560, 0, 1560, 30)

pgn_t1 = Unit('p1', 750, 20, 125, 0)
pgn_t2 = Unit('p2', 12000, 500, 2198, 0)
pgn_t3 = Unit('p3', 57600, 2500, 6824, 0)

all_acus = [acu_t1, acu_t2, acu_t3]
all_facs = [fac_navy_t1, fac_navy_t2_hq, fac_navy_t3_hq, fac_navy_t2_sl, fac_navy_t3_sl]
all_engy = [engy_t1, engy_t2, engy_t3]
all_pgn = [pgn_t1, pgn_t2, pgn_t3]
all_units = all_acus + all_facs + all_engy + all_pgn

name_to_unit = dict([(u.name, u) for u in all_units])

"""
for u in all_pgn:
    print('%5s \t e/e: %.4f \t e/bt %.4f' % (u.name, u.e_per_e(), u.e_per_bt()))

for u in all_acus+all_engy+all_facs:
    print('%5s \t bp/e: %.4f \t bp/bt %.4f' % (u.name, u.bp_per_e(), u.bp_per_bt()))

for k, v in name_to_unit.items():
    print(k, v.name)

notes:
t1 fac wins in bp/bt over everything else, can only build t1 engies though
all engies are equal in bp/e, better than anything else, slave facs can therefore be ignored
upgrading acu to t3 is better in terms of bp/bt, fac in terms of bp/e
higher tech pgns are more efficient in both e/e and e/bt -> no going back in tech
"""


class Run:
    def __init__(self, transition_with_acu=True, e_stop=30000):
        self.units = [acu_t1, fac_navy_t1]
        self.e_prod = sum([u.e_prod for u in self.units])
        self.fac_bp = sum([u.bp for u in self.units if u.is_fac()])
        self.engy_bp = sum([u.bp for u in self.units if not u.is_fac()])
        self.t = fac_navy_t1.bt / acu_t1.bp
        self.e = 4000 - fac_navy_t1.e_cost + self.t * acu_t1.e_prod

        self.t2_trans = [acu_t2] if transition_with_acu else [fac_navy_t2_hq, engy_t2]
        self.t3_trans = [acu_t3] if transition_with_acu else [fac_navy_t3_hq, engy_t3]
        self.queues = [[], [], []]

        self.t1_options = [fac_navy_t1, engy_t1, pgn_t1]
        self.t2_options = [fac_navy_t1, engy_t1, pgn_t2]
        self.t3_options = [fac_navy_t1, engy_t1, pgn_t3]
        self.options = [self.t1_options, self.t2_options, self.t3_options]

        self.e_stop = e_stop
        self.transition_with_acu = transition_with_acu

    def mutated_copy(self, max_mutations=3):
        run = Run(transition_with_acu=self.transition_with_acu, e_stop=self.e_stop)
        run.init_from(*self.get_queues())
        for __ in range(random.randint(1, max_mutations)):
            run.mutate()
        return run

    def mutate(self):
        """ random insertion or deletion of a queue element """
        while True:
            r_q = random.randint(0, len(self.queues)-1)
            q = self.queues[r_q]
            r_i = random.randint(0, len(q))
            r_o = random.randint(0, len(self.options[r_q]))
            # insert when rolling one of the available options
            if r_o < len(self.options[r_q]):
                self.queues[r_q].insert(r_i, self.options[r_q][r_o])
                return
            # remove something, if possible
            elif r_i < len(q):
                self.queues[r_q].pop(r_i)
                return

    def init_from(self, t1q, t2t, t2q, t3t, t3q):
        """ each is a list of strings """
        self.queues = [[], [], []]
        self.queues[0] = [name_to_unit.get(s) for s in t1q]
        self.t2_trans = [name_to_unit.get(s) for s in t2t]
        self.queues[1] = [name_to_unit.get(s) for s in t2q]
        self.t3_trans = [name_to_unit.get(s) for s in t3t]
        self.queues[2] = [name_to_unit.get(s) for s in t3q]

    def print_stats(self):
        print('t: %.2f \t e_prod: %d \t e_storage: %.1f \t f_bp: %d \t e_bp: %d'
              % (self.t, self.e_prod, self.e, self.fac_bp, self.engy_bp))

    def get_queues(self):
        return [u.name for u in self.queues[0]],\
               [u.name for u in self.t2_trans],\
               [u.name for u in self.queues[1]],\
               [u.name for u in self.t3_trans],\
               [u.name for u in self.queues[2]]

    def print_queue(self):
        print(', '.join(str(q) for q in self.get_queues()))

    def evaluate(self, verbose=False):
        def steps(units: [Unit]):
            for u in units:
                # figure out time/e to build the unit
                bp = (self.fac_bp + self.engy_bp) if u.is_engy() else self.engy_bp  # which bp for current unit
                td_b = u.bt / bp                                                    # time by bp alone
                e_later = self.e + td_b * self.e_prod                               # energy in storage by then
                td_e = (u.e_cost - e_later) / self.e_prod                           # remaining time for e
                td = td_b + max(0, td_e)                                            # total required time to build

                # update values
                self.e = self.e + self.e_prod * td - u.e_cost   # e storage
                self.e_prod += u.e_prod                         # e production
                self.fac_bp += u.bp if u.is_fac() else 0        # fac bp
                self.engy_bp += u.bp if not u.is_fac() else 0   # engy bp
                self.t += td                                    # passed time

                if self.e_prod >= self.e_stop:
                    return True
            return False

        for unit_groups in [self.queues[0], self.t2_trans, self.queues[1], self.t3_trans, self.queues[2]]:
            done = steps(unit_groups)
            if verbose:
                self.print_stats()
            if done:
                return
        while not steps([pgn_t3]):
            pass


class Climber:
    def __init__(self, init_run: Run, pool_size=10, max_mutations=5):
        init_run.evaluate()
        self.all = [init_run]
        self.pool_size = pool_size
        self.max_mutations = max_mutations
        self.score = init_run.t

    def step(self):
        # insert runs and evaluate
        for _ in range(self.pool_size-1):
            run = self.all[0].mutated_copy(max_mutations=self.max_mutations)
            run.evaluate(verbose=False)
            self.all.append(run)
        # keep only the best alive
        self.all = sorted(self.all, key=lambda u: u.t)

        # for a in self.all:
        #     print(a.t)

        self.all = self.all[0:1]
        s = self.all[0].t
        if s < self.score:
            print('-'*50)
            self.all[0].print_stats()
            self.all[0].print_queue()
            self.score = s


t0 = time.time()
r = Run(transition_with_acu=False)
# r.init_from(['e1', 'p1', 'p1'], ['a2'], [], ['a3'], [])
r.print_queue()

print('-'*100)
print('starting search')
print('-'*100)

climber = Climber(r, pool_size=5)
for i in range(10000):
    climber.step()

print('\n')
print('search time: %.2f' % (time.time() - t0))