monstercount

#!/usr/bin/env python

from struct import unpack
from collections import Counter
from termcolor import colored
from os.path import basename, splitext
from zipfile import ZipFile
import argparse

EXCESSIVE_MONSTER_COUNT = 300

# Thing Types | https://doomwiki.org/wiki/Thing_types#Monsters
MONSTER_TYPES = { 3004:"tr", 9:"sg", 65:"cg", 3001:"im",
                  3002:"dm", 58:"sp", 3006:"ls", 3005:"ca",
                  69:"hk", 3003:"br", 68:"ar", 71:"pe",
                  66:"rv", 67:"mn", 64:"vi", 7:"sd",
                  16:"cy", 84:"ss", 72:"ke" }

# Thing flags | https://doomwiki.org/wiki/Thing#Flags
FLAG_HNTR = 1<<0 # 0x01
FLAG_HMP  = 1<<1 # 0x02
FLAG_UV   = 1<<2 # 0x04
FLAG_DEAF = 1<<3 # 0x08
FLAG_MP   = 1<<4 # 0x10

def get_monsters(wad, args):
    file_id, num_lumps, directory_offset = unpack('4sii', wad.read(12))

    if file_id not in (b'PWAD', b'IWAD'):
        raise ValueError("Not a valid WAD file")

    monster_data = []
    map_name = ""

    if args.skill in (1, 2):
        skill_flag = FLAG_HNTR
    elif args.skill == 3:
        skill_flag = FLAG_HMP
    elif args.skill in (4, 5):
        skill_flag = FLAG_UV
    flag_mask = skill_flag | FLAG_MP

    for i in range(num_lumps):
        wad.seek(directory_offset + 16 * i)
        lump_offset, lump_size, lump_name = unpack('ii8s', wad.read(16))
        lump_name = lump_name.split(b'\0', 1)[0].decode()

        if lump_name == 'THINGS':
            wad.seek(lump_offset)
            map_monsters = []
            while lump_size >= 10:
                x, y, angle, type_id, flags = unpack('hhHHH', wad.read(10))
                if type_id in MONSTER_TYPES and (flags & flag_mask) == skill_flag:
                    map_monsters.append(type_id)
                lump_size -= 10
            monster_data.append((map_name, map_monsters))
        else:
            map_name = lump_name

    return monster_data

def tabulate(monster_data):
    monster_counts = {}
    for map_name, map_monsters in monster_data:
        monster_counts.setdefault(map_name, Counter()).update(map_monsters)

    present_monsters = {mon for counter in monster_counts.values() for mon in counter}
    headers = [MONSTER_TYPES[key] for key in MONSTER_TYPES if key in present_monsters]

    table = []
    for row in sorted(monster_counts):
        row_totals = [monster_counts[row].get(key, 0) for key in MONSTER_TYPES if MONSTER_TYPES[key] in headers]
        table.append([row] + [str(count) for count in row_totals] + [str(sum(row_totals))])

    column_totals = [sum(int(row[i]) for row in table) for i in range(1, len(table[0]))]
    table.append(["TOTAL"] + [str(count) for count in column_totals])

    table.insert(0, [""] + headers + ["TOTAL"])

    return table

def display(tabulated_data, args):
    formatted_table = [["" if cell == '0' else cell for cell in row] for row in tabulated_data]

    formatted_table = formatted_table[:-1]

    colwidths = [max(len(x) for x in col) for col in zip(*formatted_table)]
    formatted_table = [[x.rjust(w) for x,w in zip(row, colwidths)] for row in formatted_table]

    for row in formatted_table[:-1]:
        last_col_idx = len(row) - 1
        value = row[last_col_idx]
        if value.strip().isdigit() and int(value) > EXCESSIVE_MONSTER_COUNT:
            row[last_col_idx] = colored(value, 'red')

    for row in formatted_table:
        print(" ".join(row))

    if len(formatted_table) > 24:
        print(" ".join(formatted_table[0]))

    wad_count = len(tabulated_data) - 2
    populated_count = sum(1 for row in tabulated_data[1:-1] if int(row[-1]) > 0)
    total_monsters = int(tabulated_data[-1][-1])
    avg_monsters = round(total_monsters / populated_count) if populated_count else 0
    print(f"WAD count: {populated_count} Avg monsters: {avg_monsters}")

def process_wad(wad, args):
    monster_data = get_monsters(wad, args)
    if not monster_data:
        print("No maps found.")
        return
    tabulated_data = tabulate(monster_data)
    display(tabulated_data, args)

def main(args):
    for file in args.files:
        if file.lower().endswith('.wad'):
            print(f"Opening {file} with -skill {args.skill}...")
            with open(file, 'rb') as wad:
                process_wad(wad, args)
        elif file.lower().endswith('.zip'):
            with ZipFile(file) as zf:
                zip_name = splitext(basename(file))[0].lower()
                wad_name = f"{zip_name}.wad"
                namelist = zf.namelist()
                fn = (next((n for n in namelist if n.lower() == wad_name), None)
                   or next((n for n in namelist if n.lower().endswith('.wad')), None))
                if fn is not None:
                    print(f"Opening {fn} inside {file} with -skill {args.skill}...")
                    with zf.open(fn, 'r') as wad:
                        process_wad(wad, args)
                else:
                    print("No .wad found inside archive")
        else:
            print("Unsupported file type. This script only accepts .wad and .zip files.")

def parse_arguments():
    parser = argparse.ArgumentParser(description="Parses Doom wad files for monster counts")
    parser.add_argument(
        '-skill',
        type=int,
        choices=range(1, 6),
        default=4,
        help="skill level (1-5), 1=ITYTD, 2=HNTR, 3=HMP, 4=UV, 5=Nightmare (defaults to (4) UV)"
    )
    parser.add_argument(
        'files',
        nargs='*',
        help=".wad or .zip files to process"
    )
    return parser.parse_args()

if __name__=='__main__':
    args = parse_arguments()
    if args.files:
        main(args)
    else:
        print('Please provide at least one file to process.')