4-Esper Challenge generator

#!/usr/bin/env python3
#
# Copyright (C) 2019-2022 Síle Ekaterin Liszka
#
# Permission to use, copy, modify, and/or distribute this software for
# any purpose with or without fee is hereby granted, provided that the
# above copyright notice and this permission notice appear in all
# copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
# WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
# AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR
# CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
# OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#
# This project was made possible by:
# * Lenophis, who told me where the 'give this esper' bytes are
# * ff6hacking.com, which hosts the data tables I used to pick from.
# * CanOfWorms, for figuring out the GBA offsets.
#
# ChangeLog:
# 0.3:
# - GBA support! Purportedly works for J, U, *and* E versions.
# 0.2:
# - Preliminary work toward supporting GBA ROMs. This needs data I
#   don't have to complete, however, but the cleanup it required was
#   useful.
# - Reworked learning rates. Rates should be overall lower now.
# - Majority of output moved to a spoilers file.

import sys
import io

from random import randint, seed, shuffle, random, choice, sample
from enum import IntEnum
from copy import copy

class Bonus(IntEnum):
    HP_10 = 0x00
    HP_30 = 0x01
    HP_50 = 0x02
    MP_10 = 0x03
    MP_30 = 0x04
    MP_50 = 0x05
    HP100 = 0x06
    LV_30 = 0x07
    LV_50 = 0x08
    Str_1 = 0x09
    Str_2 = 0x0A
    Spd_1 = 0x0B
    Spd_2 = 0x0C
    Vit_1 = 0x0D
    Vit_2 = 0x0E
    Mag_1 = 0x0F
    Mag_2 = 0x10
    Nil   = 0xFF

bonuses = {
    Bonus.HP_10: 'HP +10%',
    Bonus.HP_30: 'HP +30%',
    Bonus.HP_50: 'HP +50%',
    Bonus.MP_10: 'MP +10%',
    Bonus.MP_30: 'MP +30%',
    Bonus.MP_50: 'MP +50%',
    Bonus.HP100: 'HP +100%',
    Bonus.Str_1: 'Strength +1',
    Bonus.Str_2: 'Strength +2',
    Bonus.Spd_1: 'Speed +1',
    Bonus.Spd_2: 'Speed +2',
    Bonus.Mag_1: 'Magic +1',
    Bonus.Mag_2: 'Magic +2',
}

class Spell(IntEnum):
    Fire       = 0x00
    Blizzard   = 0x01
    Thunder    = 0x02
    Poison     = 0x03
    Drain      = 0x04
    Firera     = 0x05
    Blizzardra = 0x06
    Thunderra  = 0x07
    Bio        = 0x08
    Firega     = 0x09
    Blizzardga = 0x0A
    Thunderga  = 0x0B
    Break      = 0x0C
    Death      = 0x0D
    Holy       = 0x0E
    Flare      = 0x0F
    Gravity    = 0x10
    Gravityga  = 0x11
    Dezone     = 0x12
    Meteor     = 0x13
    Ultima     = 0x14
    Quake      = 0x15
    Tornado    = 0x16
    Meltdown   = 0x17
    Library    = 0x18
    Slow       = 0x19
    Rasp       = 0x1A
    Silence    = 0x1B
    Protect    = 0x1C
    Sleep      = 0x1D
    Confuse    = 0x1E
    Haste      = 0x1F
    Stop       = 0x20
    Berserk    = 0x21
    Float      = 0x22
    Kappa      = 0x23
    Reflect    = 0x24
    Shell      = 0x25
    Vanish     = 0x26
    Hastega    = 0x27
    Slowga     = 0x28
    Osmose     = 0x29
    Warp       = 0x2A
    Quick      = 0x2B
    Dispel     = 0x2C
    Cure       = 0x2D
    Curera     = 0x2E
    Curega     = 0x2F
    Raise      = 0x30
    Arise      = 0x31
    NullPoison = 0x32
    NullStatus = 0x33
    Regen      = 0x34
    Reraise    = 0x35

    Flood   = 54
    Gravija = 55
    Valor   = 56

    SNES_Count = 54
    GBA_Count  = 57

class Esper(IntEnum):
    Ramuh         =  0
    Ifrit         =  1
    Shiva         =  2
    Siren         =  3
    Jormungand    =  4
    Catoblepas    =  5
    Maduin        =  6
    Bismarck      =  7
    CatSidhe      =  8
    Quetzali      =  9
    Valigarmander = 10
    Odin          = 11
    Raiden        = 12
    Bahamut       = 13
    Alexander     = 14
    Jihad         = 15
    Ragnarok      = 16
    QiLin         = 17
    ZonaSeeker    = 18
    Carbunkle     = 19
    Phantom       = 20
    Seraphim      = 21
    Golem         = 22
    Unicorn       = 23
    Fenrir        = 24
    Lakshmi       = 25
    Phoenix       = 26

    Leviathan = 27
    Cactender = 28
    Diabolus  = 29
    Gilgamesh = 30

    SNES_Count = 27
    GBA_Count = 31

def roll_espers(max):
    return [Esper(i) for i in sample(list(range(max)), 4)]

spell_rate = [1, 1, 1, 2, 2, 3, 4, 3, 2, 2, 1, 1, 1]

def roll_spells(count, max, forbidden=[]):
    spells = []
    spell_base = []

    def add_spell(spell):
        spell_base.append(spell)
        spells.append([spell, choice(spell_rate) * 5])

    if randint(1, 100) < 50: add_spell(Spell.Curera)
    else:                    add_spell(Spell.Curega)

    num = randint(1, 100)
    if num < 34:
        add_spell(Spell.Fire)
        num2 = randint(1, 100)
        if num2 < 50: add_spell(Spell.Blizzardra)
        else:         add_spell(Spell.Thunderra)
    elif num < 67:
        add_spell(Spell.Blizzard)
        num2 = randint(1, 100)
        if num2 < 50: add_spell(Spell.Firera)
        else:         add_spell(Spell.Thunderra)
    else:
        add_spell(Spell.Thunder)
        num2 = randint(1, 100)
        if num2 < 50: add_spell(Spell.Firera)
        else:         add_spell(Spell.Blizzardra)

    add_spell(Spell.Protect)
    add_spell(Spell.Shell)
    add_spell(Spell.Haste)

    num = randint(1, 100)
    if num < 34:   add_spell(Spell.Raise)
    elif num < 67: add_spell(Spell.Arise)

    n = count - len(spells)
    while True:
        if n < 0:
            break

        spell = Spell(randint(0, max))

        if spell not in spell_base and spell not in forbidden:
            add_spell(spell)
            n -= 1

    return spells

def generate(rng_seed=None, mode='snes'):
    seed(rng_seed)

    if mode == 'snes':
        esper_max = Esper.SNES_Count - 1
        spell_max = Spell.SNES_Count - 1
    elif mode == 'gba':
        esper_max = Esper.GBA_Count - 1
        spell_max = Spell.GBA_Count - 1

    espers = roll_espers(esper_max)
    statup = [Bonus(i) for i in sample(list(bonuses.keys()), 4)]

    spells = roll_spells(20, spell_max)
    shuffle(spells)

    return {
        'espers': espers,
        'level+': statup,
        'spells': spells,
        'seed':   rng_seed
    }

remove = [
    0xA55DE, 0xB4DE3, 0xB52C9, 0xB5452,
    0xA5A2C, 0xB5B77, 0xB9A7C, 0xC1ED1,
    0xC1F85, 0xC2044, 0xC316C, 0xC37A8,
    0xC4ADA, 0xC5E08, 0xC5E08, 0xC79D9,
    0xC79E9, 0xCD78F
]
def write_rom(inname, outname, mode, data):
    print('Writing data for seed {} to {}...'.format(data['seed'], realname))

    offset = 0
    if mode == 'gba':
        offset = 0x70000
        spello = 0x62A480
        remove.extend((0xD2B52, 0xD2BD4, 0xD11ED, 0xD27A5))
    else:
        spello = 0x186e00

    rom = io.BytesIO()

    with open(inname, mode='rb') as f:
        rom.write(f.read())

    # erase all the slots we're removing
    # Bahamut
    rom.seek(0xA00DE + offset, 0)
    rom.write(b'\xFD' * 5)
    # sorcerous research facility
    rom.seek(0xC7C6D + offset, 0)
    rom.write(b'\xFD' * 11)
    # everybody else
    for i in remove:
        rom.seek(i + offset, 0)
        rom.write(b'\xFD' * 2)

    # these four are the espers given when you reach Tina in Zozo
    rom.seek(0xAA824 + offset, 0)
    rom.write(bytes([data['espers'][0] + 0x36]))
    rom.seek(0xAAC9C + offset, 0)
    rom.write(bytes([data['espers'][1] + 0x36]))
    rom.seek(0xAACAB + offset, 0)
    rom.write(bytes([data['espers'][2] + 0x36]))
    rom.seek(0xAACBA + offset, 0)
    rom.write(bytes([data['espers'][3] + 0x36]))

    # write esper data for the espers we're giving out
    for i in range(4):
        rom.seek(spello + (data['espers'][i] * 11), 0)
        for j in range(5):
            spell = data['spells'].pop(0)
            rom.write(bytes([spell[1], spell[0]]))
        rom.write(bytes([data['level+'][i]]))

    rom.seek(0, 0)

    with open(outname, mode='wb') as f:
        f.write(rom.read())

if __name__ == '__main__':
    from sys import argv
    from time import time

    now = str(time()).split('.')[0]

    if len(argv) == 3:
        rng_seed = argv[2]
    elif len(argv) < 2 or len(argv) > 3:
        print('Usage: {} <filename> [seed]'.format(argv[0]))
    else:
        rng_seed = now

    filename = argv[1]

    if filename.endswith('.gba'):
        mode = 'gba'
    else:
        mode = 'snes'
    print('Mode: {}'.format(mode))

    print('Seed: {}'.format(rng_seed))
    data = generate(rng_seed, mode)

    filefrags = filename.split('.')
    ext = filefrags[-1]
    filefrags[-1] = str(data['seed'])
    filefrags.append(ext)
    realname = '.'.join(filefrags)
    filefrags[-1] = 'txt'
    spoilers = '.'.join(filefrags)

    espers = []
    for i in range(4):
        esper = str(data['espers'][i]).split('.')[1]
        bonus = bonuses[data['level+'][i]]
        espers.append(esper + '(' + bonus + ')')

    spells = data['spells']

    with open(spoilers, mode='wt', encoding='utf8') as f:
        f.write('Seed: {}\n'.format(rng_seed))
        for i in range(4):
            f.write('{}\n'.format(espers[i]))

            for j in range(5):
                spell = data['spells'][(i * 5) + j]
                f.write('  {} ×{}\n'.format(str(spell[0]).split('.')[1], spell[1]))

    write_rom(filename, realname, mode, data)