import re
from lib import get_data, ints
from collections import defaultdict
from copy import deepcopy

r = re.compile(r"(immune|weak) to ([\w ,]+)")
r_strength = re.compile(r"(\w+) damage")

data = """Immune System:
17 units each with 5390 hit points (weak to radiation, bludgeoning) with an attack that does 4507 fire damage at initiative 2
989 units each with 1274 hit points (immune to fire; weak to bludgeoning, slashing) with an attack that does 25 slashing damage at initiative 3

Infection:
801 units each with 4706 hit points (weak to radiation) with an attack that does 116 bludgeoning damage at initiative 1
4485 units each with 2961 hit points (immune to radiation; weak to fire, cold) with an attack that does 12 slashing damage at initiative 4
 """
data = get_data(__file__)

ARMY, ID, UNITS, HITS, DAMG, INIT = 0, 1, 2, 3, 4, 5

groups = []
unit_meta = {}

id = 1
for a, p in enumerate(data.split("\n\n")):
    for line in p.splitlines():
        xs = ints(line)
        if len(xs) == 0:
            continue
        assert len(xs) == 4
        xs = [a, id] + xs

        (attack,) = r_strength.findall(line)
        defense = defaultdict(list)
        for attr, values in r.findall(line):
            defense[attr] = values.split(", ")

        unit_meta[id] = attack, defense
        groups.append(xs)
        id += 1


def effective_power(group):
    return (group[UNITS] * group[DAMG], group[INIT])


def damage(g1, g2):
    """Damage g1 does to g2"""
    g1_id, g2_id = g1[ID], g2[ID]
    g1_attack_type = unit_meta[g1_id][0]
    g2_defenses = unit_meta[g2_id][1]
    g2_immune = True if g1_attack_type in g2_defenses["immune"] else False
    g2_weak = True if g1_attack_type in g2_defenses["weak"] else False

    assert (g2_immune and g2_weak) != True

    damage = g1[UNITS] * g1[DAMG]
    if g2_immune:
        damage = 0
    elif g2_weak:
        damage *= 2
    return damage


groups_orig = deepcopy(groups)
while True:
    groups = sorted(groups, key=effective_power, reverse=True)
    attacks = {}
    attacked = set()
    for g in groups:
        target_max_stats = 0, 0, 0
        target = None
        for t in groups:
            if g[ARMY] == t[ARMY]:
                continue
            if g[ID] == t[ID]:
                continue
            if t[ID] in attacked:
                continue
            target_stats = damage(g, t), effective_power(t)[0], t[INIT]
            if target_stats[0] == 0:
                continue
            if target_stats > target_max_stats:
                target_max_stats = target_stats
                target = t

        if target is not None:
            attacked.add(target[ID])
            attacks[g[ID]] = target

    groups = sorted(groups, key=lambda g: g[INIT], reverse=True)
    for g in groups:
        if g[ID] not in attacks:
            continue
        if g[UNITS] <= 0:
            continue

        target = attacks[g[ID]]
        cdamage = damage(g, target)
        units_dead = cdamage // target[HITS]
        units_dead = units_dead if units_dead <= target[UNITS] else target[UNITS]
        # print("immune" if g[ARMY] == 0 else "infection", g[ID], "attacks", target[ID], "dead", units_dead)
        target[UNITS] -= units_dead

    groups = [g for g in groups if g[UNITS] > 0]
    armies = set()
    for g in groups:
        armies.add(g[ARMY])
    assert len(armies) >= 1
    if len(armies) == 1:
        break


t = sum(g[UNITS] for g in groups)
print(t)


for boost in range(0, 1800):
    groups = deepcopy(groups_orig)
    for g in groups:
        if g[ARMY] == 0:
            g[DAMG] += boost

    while True:
        groups = sorted(groups, key=effective_power, reverse=True)
        attacks = {}
        attacked = set()
        for g in groups:
            target_max_stats = 0, 0, 0
            target = None
            for t in groups:
                if g[ARMY] == t[ARMY]:
                    continue
                if g[ID] == t[ID]:
                    continue
                if t[ID] in attacked:
                    continue
                target_stats = damage(g, t), effective_power(t)[0], t[INIT]
                if target_stats[0] == 0:
                    continue
                if target_stats > target_max_stats:
                    target_max_stats = target_stats
                    target = t

            if target is not None:
                attacked.add(target[ID])
                attacks[g[ID]] = target

        groups = sorted(groups, key=lambda g: g[INIT], reverse=True)
        max_dead = 0
        for g in groups:
            if g[ID] not in attacks:
                continue
            if g[UNITS] <= 0:
                continue

            target = attacks[g[ID]]
            cdamage = damage(g, target)
            units_dead = cdamage // target[HITS]
            units_dead = units_dead if units_dead <= target[UNITS] else target[UNITS]
            max_dead = max(units_dead, max_dead)
            # print("immune" if g[ARMY] == 0 else "infection", g[ID], "attacks", target[ID], "dead", units_dead)
            target[UNITS] -= units_dead

        if max_dead == 0:
            break

        groups = [g for g in groups if g[UNITS] > 0]
        armies = set()
        for g in groups:
            armies.add(g[ARMY])
        assert len(armies) >= 1
        if len(armies) == 1:
            army = groups[0][ARMY]
            t = sum(g[UNITS] for g in groups)
            if army == 0:
                print(t)
                exit(0)
            break