builder.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import getopt
import yaml
import csv
import tarfile
import re
import configparser
from string import Template

from converter import convert_demo_to_libtas


N_LEVELS_PATH = "volume/n_levels"

def usage():
    print(f"Usage: {sys.argv[0]} [-e END_EPISODE|-s START_EPISODE|-r|-h]")
    print("-i LEVEL: build infos for given level (for display_infos.lua)")
    print("-r: build RTA scores")
    print("-h: print this help")


def start_episode(col, row):
    with open("tas/episode_click_coordinates.yml", "r",
              encoding="utf-8") as f:
        coord = yaml.safe_load(f)
    # On the first frame, press n ("K6e") while moving the mouse to
    # the right coordinates. On the second frame, click on right
    # coordinates
    return f"|K6e|M{coord["column"][col]}:{coord["row"][row]}:A:.....:0|\n|M{coord["column"][col]}:{coord["row"][row]}:A:1....:0|\n|\n"


def get_rerecords_number(level_name):
    archive_path = f"{N_LEVELS_PATH}/{level_name}.ltm"

    with tarfile.open(archive_path, "r:gz") as tar:
        try:
            config_member = tar.getmember("config.ini")
        except KeyError:
            raise FileNotFoundError("config.ini not found in archive")

        config_file = tar.extractfile(config_member)
        if config_file is None:
            raise ValueError("Could not extract config.ini")

        content = config_file.read().decode("utf-8")

    match = re.search(r"^rerecord_count\s*=\s*(\d+)", content, re.MULTILINE)
    if not match:
        raise KeyError("rerecord_count not found in config.ini")

    return int(match.group(1))


def build_libtas_input(begin_episode=0, end_episode=99, begin_level=0, end_level=4, rta=False, score_type="Speedrun"):
    nb_frames = 0
    res = ""
    markers = {}
    nb_markers = 0
    lua_infos = ""
    initial_wait_frames = 7
    for i in range(initial_wait_frames):
        res += "|\n"
        nb_frames += 1
    start_col = int(begin_episode/10)
    start_row = begin_episode % 10
    res += start_episode(start_col, start_row)
    nb_frames += 3
    if rta:
        level_data_file="tas/level_data_rta.yml"
    else:
        level_data_file="tas/level_data.yml"
        level_data_file_rta = "tas/level_data_rta.yml"
        with open(level_data_file_rta, "r", encoding="utf-8") as f:
            data_rta = yaml.safe_load(f)
    with open(level_data_file, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    with open("tas/loading_times.yml", "r", encoding="utf-8") as f:
        loading_times = yaml.safe_load(f)
    output_ghost_file = open("volume/n_recomp_rta_speedrun_ghost.csv", "w", encoding="utf-8")
    output_ghost = csv.writer(output_ghost_file)

    for level_name, level_data in data.items():
        episode = int(level_name.split("-")[0])
        level = int(level_name.split("-")[1])
        if episode < begin_episode or (episode == begin_episode and level < begin_level):
            # skip to the beginning
            continue
        elif episode > end_episode or (episode == end_episode and level > end_level):
            # we reached the end
            break
        ghost_file = f"volume/ghosts/{level_name}.csv"
        with open(ghost_file, "r", encoding="utf-8") as f:
            reader = csv.reader(f)
            ghost_data = list(reader)
        if loading_times[level_name] is None:
            loading_time = 58  # fill missing ones so as to be able to calculate them with script
        else:
            loading_time = int(loading_times[level_name])
        for i in range(loading_time):
            res += "|\n"
            nb_frames += 1
        res += "|K20|\n"  # space
        nb_frames += 1
        if score_type in level_data and "demo" in level_data[score_type]:
            nb_markers += 1
            markers[f"{nb_markers}\\frame"] = nb_frames
            markers[f"{nb_markers}\\text"] = level_name
            demo_str = level_data[score_type]["demo"]
            rerecords = get_rerecords_number(level_name)
            libtas_input, nb_frames_demo = convert_demo_to_libtas(demo_str)
            if rta:
                try:
                    date = level_data[score_type]['time'].strftime("%Y-%m-%d")
                except AttributeError:
                    # for the kryX-orange 25-2 case, which is not a date
                    date = level_data[score_type]['time']
                lua_infos += f"    {{{nb_frames}, {nb_frames+nb_frames_demo}, \"{level_data[score_type]['authors']}, {level_data[score_type]['time']}, {date}\"}},\n"
            else:
                diff_with_rta = int(data_rta[level_name][score_type]['time']) - int(level_data[score_type]['time'].replace(' f', ''))
                lua_infos += f"    {{{nb_frames}, {nb_frames+nb_frames_demo}, \"TAS: {level_data[score_type]['authors']}, {level_data[score_type]['time']}\",\"0th: {data_rta[level_name][score_type]['authors']}, {data_rta[level_name][score_type]['time']} f\", \"- {diff_with_rta} f\", {rerecords}}},\n"
            res += libtas_input
            for i in range(nb_frames_demo):
                output_ghost.writerow([nb_frames + i] + ghost_data[i][1:])
            nb_frames += nb_frames_demo
            res += "|K20|\n"  # space
            nb_frames += 1
        if int(level_name.split("-")[1]) == 4 and end_level == 4:
            # pass end of episode screen
            res += "|\n|K20|\n"
            nb_frames += 2
            if episode % 10 == 9 and episode < end_episode:
                res += "|\n"
                res += start_episode(int((episode + 1) / 10), 0)
                nb_frames += 4
    markers["size"] = nb_markers
    output_ghost_file.close()
    return res, nb_frames, markers, lua_infos


def parse_args():
    starting_episode = 0
    end_episode = None
    rta = False
    infos_level = None
    try:
        opts, args = getopt.getopt(sys.argv[1:], 'e:hi:rs:')
    except getopt.GetoptError as err:
        print("Error: ", str(err))
        sys.exit(1)

    for o, arg in opts:
        if o == '-e':
            end_episode = int(arg)
        if o == '-r':
            rta = True
        if o == '-s':
            starting_episode = int(arg)
        if o == '-i':
            level_name = arg
            starting_episode = int(level_name.split("-")[0])
            infos_level = int(level_name.split("-")[1])
        if o == '-h':
            usage()
    if end_episode is None:
        end_episode = starting_episode
    return starting_episode, end_episode, rta, infos_level


if __name__ == "__main__":
    starting_episode, end_episode, rta, infos_level = parse_args()
    if infos_level is not None:
        begin_level = infos_level
        end_level = infos_level
    else:
        begin_level = 0
        end_level = 4
    libtas_input, nb_frames, markers, lua_infos = build_libtas_input(starting_episode, end_episode, begin_level=begin_level, end_level=end_level, rta=rta, score_type="Speedrun")
    with open("infos.lua.template") as f:
        template = Template(f.read())
    lua_script_filled = template.substitute(infos=lua_infos)
    with open("volume/lua/data/infos.lua", "w") as f:
        print(lua_script_filled, file=f)
    if not infos_level:
        with open("extract/inputs", "w") as f:
            print(libtas_input, file=f)
        config = configparser.ConfigParser(strict=False, delimiters=('='), interpolation=None)
        config.read("extract/editor.ini")
        with open("extract/editor.ini", "w") as f:
            config["markers"] = markers
            config.write(f, space_around_delimiters=False)