metroidgen/metroidgen.py

# Metroid (NES) Random Password Generator
# Author: Noëlle Anthony
# License: MIT
# Date: October 2019
# This uses http://games.technoplaza.net/mpg/password.txt as a basis for its password algorithm

import random, sys

class MetroidState:
    """ Stores the game state
    """
    def __init__(self):
        # Alphabet is 64 characters - 6 bits per character
        self.alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz?-"
        # The password has different flags for "item available for pickup" and
        # "Samus has the item". I'm keeping them separate, but when the generator
        # selects an item as "picked up", that means Samus has it and it's not available.
        self.itemsCollected = {
            "Maru Mari": False,
            "Bombs": False,
            "Long Beam": False,
            "Ice Beam": False,
            "Wave Beam": False,
            "High Jump Boots": False,
            "Varia": False,
            "Screw Attack": False
        }
        self.samusHas = {
            "Maru Mari": False,
            "Bombs": False,
            "Long Beam": False,
            "Ice Beam": False,
            "Wave Beam": False,
            "High Jump Boots": False,
            "Varia": False,
            "Screw Attack": False
        }
        # Missile tanks are listed in the order in which they appear in the password,
        # NOT in zone order or in any reasonable collection order.
        self.missileTanks = {
            1: False,
            2: False,
            3: False,
            4: False,
            5: False,
            6: False,
            7: False,
            8: False,
            9: False,
            10: False,
            11: False,
            12: False,
            13: False,
            14: False,
            15: False,
            16: False,
            17: False,
            18: False,
            19: False,
            20: False,
            21: False
        }
        # Likewise energy tanks are listed in password order.
        self.energyTanks = {
            1: False,
            2: False,
            3: False,
            4: False,
            5: False,
            6: False,
            7: False,
            8: False
        }
        # This may be left-to-right (Samus approaches from the right). I haven't checked.
        self.zebetitesDestroyed = {
            1: False,
            2: False,
            3: False,
            4: False,
            5: False
        }
        # I'm not sure why I decided to segregate these by zone, except that that's how
        # truepeacein.space does it.
        self.doors = {
            "Brinstar": {
                1: False,
                2: False,
                3: False,
                4: False,
                5: False
            }, "Norfair": {
                1: False,
                2: False,
                3: False,
                4: False
            }, "Kraid": {
                1: False,
                2: False,
                3: False,
                4: False,
                5: False
            }, "Ridley": {
                1: False,
                2: False
            }, "Tourian": {
                1: False,
                2: False,
                3: False
            }
        }
        # The next three are self-explanatory.
        self.kraidKilled = False
        self.ridleyKilled = False
        self.motherBrainKilled = False
        # The Kraid and Ridley statues rise when Kraid and Ridley are killed, but
        # their states are stored separately in the password. It's possible to 
        # raise them without killing the bosses, granting early access to Tourian.
        self.kraidStatue = False
        self.ridleyStatue = False
        # Is Samus wearing her armor (False) or her swimsuit (True)?
        self.swimsuit = False
        # 0-255. You can have more missiles than 5*collected tanks (in fact, you 
        # can only collect 21 tanks - thus 105 missiles - but can have up to 255
        # missiles in your inventory).
        self.missileCount = 0
        # How advanced is the game clock? After 3 hours you don't get the good ending.
        self.gameAge = 0
        # There are five possible start locations: Brinstar, where you start, and
        # at the bottom of the elevator where you enter each subsequent zone.
        self.locations = ["Brinstar", "Norfair", "Kraid's Lair", "Ridley's Lair", "Tourian"]
        self.startLocation = 0
        # Arrays to store the 144 bits that compose the password
        self.bitfield = []
        self.initializeBitfield()
        self.fullbitfield = []

    def initializeBitfield(self):
        """ Set the first 128 bits of the bitfield to 0.
            The remaining 16 bits will be set later.
        """
        self.bitfield = []
        for _ in range(128):
            self.bitfield.append(0)

    def toggleItem(self, itm):
        """ Mark an item as collected or uncollected.
        """
        if itm in self.itemsCollected.keys():
            self.itemsCollected[itm] = not self.itemsCollected[itm]
            self.samusHas[itm] = not self.samusHas[itm]
        else:
            print("Couldn't find item: {}".format(str(itm)))
    
    def toggleMissileTank(self, num):
        """ Mark a missile tank as collected or uncollected.
        """
        try:
            num = int(num)
        except:
            print("{} is not a number".format(num))
            return
        if num in self.missileTanks.keys():
            self.missileTanks[num] = not self.missileTanks[num]
        else:
            print("Couldn't find missile tank: {}".format(num))
    
    def toggleEnergyTank(self, num):
        """ Mark an energy tank as collected or uncollected.
        """
        try:
            num = int(num)
        except:
            print("{} is not a number".format(num))
            return
        if num in self.energyTanks.keys():
            self.energyTanks[num] = not self.energyTanks[num]
        else:
            print("Couldn't find energy tank: {}".format(num))
    
    def toggleZebetite(self, num):
        """ Mark a Zebetite stem as destroyed or intact.
        """
        try:
            num = int(num)
        except:
            print("{} is not a number".format(num))
            return
        if num in self.zebetitesDestroyed.keys():
            self.zebetitesDestroyed[num] = not self.zebetitesDestroyed[num]
        else:
            print("Couldn't find Zebetite: {}".format(num))

    def toggleKraid(self):
        """ Mark Kraid as killed or alive.
        """
        self.kraidKilled = not self.kraidKilled
        self.kraidStatue = self.kraidKilled

    def toggleKraidStatue(self):
        """ Mark Kraid's statue as raised or lowered.
            If Kraid is killed but his statue isn't raised, you can't complete the game.
        """
        self.kraidStatue = not self.kraidStatue
        if self.kraidKilled and not self.kraidStatue:
            print("WARNING: Kraid has been killed but his statue has not been raised.")

    def toggleRidley(self):
        """ Mark Ridley as killed or alive.
        """
        self.ridleyKilled = not self.ridleyKilled
        self.ridleyStatue = self.ridleyKilled

    def toggleRidleyStatue(self):
        """ Mark Ridley's statue as raised or lowered.
            If Ridley is killed but his statue isn't raised, you can't complete the game.
        """
        self.ridleyStatue = not self.ridleyStatue
        if self.ridleyKilled and not self.ridleyStatue:
            print("WARNING: Ridley has been killed but his statue has not been raised.")
    
    def toggleMotherBrain(self):
        """ Mark Mother Brain as killed or alive.
            If Mother Brain is marked as killed, the self-destruct timer won't go off 
            when you reach her room.
        """
        self.motherBrainKilled = not self.motherBrainKilled

    def toggleDoor(self, area, door):
        """ Mark a given red/yellow door as opened or locked.
        """
        try:
            area = str(area)
            door = int(door)
        except:
            print("Area must be string, door must be a positive integer")
            return
        if area in self.doors.keys() and int(door) in self.doors[area].keys():
            self.doors[area][door] = not self.doors[area][door]
        else:
            print("Couldn't find door {} in area {}".format(door, area))

    def toggleSwimsuit(self):
        """ Determine whether or not Samus is wearing her armor.
        """
        self.swimsuit = not self.swimsuit

    def newLocation(self, loc):
        """ Set a new starting location (0-4).
        """
        try:
            loc = str(loc)
        except:
            print("Location must be a string")
            return
        if loc in self.locations:
            self.startLocation = self.locations.index(loc)
        else:
            print("Couldn't find location: {}".format(loc))

    def collectedItems(self):
        """ List which items have been collected.
            Under this generator's rules, if Samus doesn't have an item,
            it's available to be picked up.
        """
        o = []
        for k,v in self.itemsCollected.items():
            if v == True:
                o.append(k)
        if len(o) == 0:
            return "None"
        else:
            return ", ".join(o)

    def collectedMissiles(self):
        """ List which missile tanks have been collected.
        """
        o = []
        for k, v in self.missileTanks.items():
            if v == True:
                o.append(k)
        if len(o) == 0:
            return "None"
        else:
            return ", ".join([str(b) for b in o])
    
    def collectedEtanks(self):
        """ List which energy tanks have been collected.
        """
        o = []
        for k, v in self.energyTanks.items():
            if v == True:
                o.append(k)
        if len(o) == 0:
            return "None"
        else:
            return ", ".join([str(b) for b in o])
    
    def killedZebetites(self):
        """ List which Zebetite stems have been destroyed.
        """
        o = []
        for k, v in self.zebetitesDestroyed.items():
            if v == True:
                o.append(k)
        if len(o) == 0:
            return "None"
        else:
            return ", ".join([str(b) for b in o])

    def killedBosses(self):
        """ List which bosses have been killed.
        """
        o = []
        if self.kraidKilled:
            o.append("Kraid")
        if self.ridleyKilled:
            o.append("Ridley")
        if self.motherBrainKilled:
            o.append("Mother Brain")
        if len(o) == 0:
            return "None"
        else:
            return ", ".join(o)
    
    def raisedStatues(self):
        """ List which statues have been raised.
        """
        o = []
        if self.kraidStatue:
            o.append("Kraid")
        if self.ridleyStatue:
            o.append("Ridley")
        if len(o) == 0:
            return "None"
        else:
            return ", ".join(o)

    def inBailey(self):
        """ Show whether Samus is in her swimsuit or not.
            'inBailey' refers to an old (false) explanation of the JUSTIN BAILEY
            password, in which a 'bailey' was English slang for a bathing suit,
            so with that password, Samus was "Just In Bailey" - i.e. in her swimsuit.
        """
        if self.swimsuit:
            return "Yes"
        else:
            return "No"        

    def openedDoors(self):
        """ List which red/yellow doors have been unlocked.
        """
        d = {"Brinstar": 0, "Norfair": 0, "Kraid": 0, "Ridley": 0, "Tourian": 0}
        o = []
        for k,v in self.doors["Brinstar"].items():
            if v == True:
                d["Brinstar"] = d["Brinstar"] + 1
        for k,v in self.doors["Norfair"].items():
            if v == True:
                d["Norfair"] = d["Norfair"] + 1
        for k,v in self.doors["Kraid"].items():
            if v == True:
                d["Kraid"] = d["Kraid"] + 1
        for k,v in self.doors["Ridley"].items():
            if v == True:
                d["Ridley"] = d["Ridley"] + 1
        for k,v in self.doors["Tourian"].items():
            if v == True:
                d["Tourian"] = d["Tourian"] + 1
        for k, v in d.items():
            o.append("{} {}".format(k, v))
        return ", ".join(o)

    def getBits(self):
        """ Return the bitfield in an easily-readable format.
        """
        o = []
        word = []
        for i in range(128):
            word.append(str(self.bitfield[i]))
            if len(word) == 8:
                o.append("".join(word))
                word = []
        o1 = " ".join(o[:8])
        o2 = " ".join(o[8:])
        return o1 + "\n" + o2

    def toString(self):
        """ Output the game state as a newline-delimited string.
        """
        ic = "Items Collected: {}".format(self.collectedItems())
        mt = "Missile Tanks Collected: {}".format(self.collectedMissiles())
        et = "Energy Tanks Collected: {}".format(self.collectedEtanks())
        zb = "Zebetites Killed: {}".format(self.killedZebetites())
        kb = "Bosses Killed: {}".format(self.killedBosses())
        rs = "Statues Raised: {}".format(self.raisedStatues())
        sw = "Swimsuit?: {}".format(self.inBailey())
        sl = "Start Location: {}".format(self.locations[self.startLocation])
        dr = "Unlocked Doors: {}".format(self.openedDoors())
        ms = "Missiles: {}".format(self.missileCount)
        pw = "Password: {}".format(self.password)
        return "\n".join([ic, mt, et, zb, kb, rs, sw, sl, dr, ms, pw])

    def randomize(self):
        """ The randomizer!
        """
        # Items
        if random.randint(0,1) == 1:
            self.toggleItem("Maru Mari")
        if random.randint(0,1) == 1:
            self.toggleItem("Bombs")
        if random.randint(0,1) == 1:
            self.toggleItem("Varia")
        if random.randint(0,1) == 1:
            self.toggleItem("High Jump Boots")
        if random.randint(0,1) == 1:
            self.toggleItem("Screw Attack")
        if random.randint(0,1) == 1:
            self.toggleItem("Long Beam")
        beam = random.randint(0,2)
        if beam == 1:
            self.toggleItem("Ice Beam")
        elif beam == 2:
            self.toggleItem("Wave Beam")
        # Missile Tanks
        for i in range(21):
            if random.randint(0,1) == 1:
                self.toggleMissileTank(i+1)
        # Energy Tanks
        for i in range(8):
            if random.randint(0,1) == 1:
                self.toggleEnergyTank(i+1)
        # Zebetites
        for i in range(5):
            if random.randint(0,1) == 1:
                self.toggleZebetite(i+1)
        # Bosses killed
        if random.randint(0,1) == 1:
            self.toggleKraid()
        if random.randint(0,1) == 1:
            self.toggleRidley()
        if random.randint(0,1) == 1:
            self.toggleMotherBrain()
        # Statues raised
        if not self.kraidKilled and random.randint(0,1) == 1:
            self.toggleKraidStatue()
        if not self.ridleyKilled and random.randint(0,1) == 1:
            self.toggleRidleyStatue()
        # Doors
        # Brinstar 5, Norfair 4, Kraid 5, Ridley 2, Tourian 3
        for i in range(5):
            if random.randint(0,1) == 1:
                self.doors["Brinstar"][i+1] = True
        for i in range(4):
            if random.randint(0,1) == 1:
                self.doors["Norfair"][i+1] = True
        for i in range(5):
            if random.randint(0,1) == 1:
                self.doors["Kraid"][i+1] = True
        for i in range(2):
            if random.randint(0,1) == 1:
                self.doors["Ridley"][i+1] = True
        for i in range(3):
            if random.randint(0,1) == 1:
                self.doors["Tourian"][i+1] = True        
        # Swimsuit
        # Samus has a 1/3 chance of spawning in her swimsuit.
        # There's no technical reason for this, just a personal choice.
        if random.randint(0,2) == 2:
            self.toggleSwimsuit()
        # Start Location
        self.startLocation = random.randint(0,4)
        self.missileCount = random.randint(0,255)

    def createBitfield(self):
        """ Create the 144-bit field from the game state 
            that will generate the password.
        """
        self.initializeBitfield()
        # Doing this in order, which is dumb and tedious but accurate.
        if self.itemsCollected["Maru Mari"]:
            self.bitfield[0] = 1
        if self.missileTanks[1]:
            self.bitfield[1] = 1
        if self.doors["Brinstar"][1]:
            self.bitfield[2] = 1
        if self.doors["Brinstar"][2]:
            self.bitfield[3] = 1
        if self.energyTanks[1]:
            self.bitfield[4] = 1
        if self.doors["Brinstar"][3]:
            self.bitfield[5] = 1
        if self.itemsCollected["Bombs"]:
            self.bitfield[6] = 1
        if self.doors["Brinstar"][4]:
            self.bitfield[7] = 1
        if self.missileTanks[2]:
            self.bitfield[8] = 1
        if self.energyTanks[2]:
            self.bitfield[9] = 1
        if self.doors["Brinstar"][5]:
            self.bitfield[10] = 1
        if self.itemsCollected["Varia"]:
            self.bitfield[11] = 1
        if self.energyTanks[3]:
            self.bitfield[12] = 1
        if self.missileTanks[3]:
            self.bitfield[13] = 1
        if self.missileTanks[4]:
            self.bitfield[14] = 1
        if self.doors["Norfair"][1]:
            self.bitfield[15] = 1
        if self.missileTanks[5]:
            self.bitfield[16] = 1
        if self.missileTanks[6]:
            self.bitfield[17] = 1
        if self.missileTanks[7]:
            self.bitfield[18] = 1
        if self.missileTanks[8]:
            self.bitfield[19] = 1
        if self.missileTanks[9]:
            self.bitfield[20] = 1
        if self.missileTanks[10]:
            self.bitfield[21] = 1
        if self.missileTanks[11]:
            self.bitfield[22] = 1
        if self.doors["Norfair"][2]:
            self.bitfield[23] = 1
        if self.itemsCollected["High Jump Boots"]:
            self.bitfield[24] = 1
        if self.doors["Norfair"][3]:
            self.bitfield[25] = 1
        if self.itemsCollected["Screw Attack"]:
            self.bitfield[26] = 1
        if self.missileTanks[12]:
            self.bitfield[27] = 1
        if self.missileTanks[13]:
            self.bitfield[28] = 1
        if self.doors["Norfair"][4]:
            self.bitfield[29] = 1
        if self.energyTanks[4]:
            self.bitfield[30] = 1
        if self.missileTanks[14]:
            self.bitfield[31] = 1
        if self.doors["Kraid"][1]:
            self.bitfield[32] = 1
        if self.missileTanks[15]:
            self.bitfield[33] = 1
        if self.missileTanks[16]:
            self.bitfield[34] = 1
        if self.doors["Kraid"][2]:
            self.bitfield[35] = 1
        if self.energyTanks[5]:
            self.bitfield[36] = 1
        if self.doors["Kraid"][3]:
            self.bitfield[37] = 1
        if self.doors["Kraid"][4]:
            self.bitfield[38] = 1
        if self.missileTanks[17]:
            self.bitfield[39] = 1
        if self.missileTanks[18]:
            self.bitfield[40] = 1
        if self.doors["Kraid"][5]:
            self.bitfield[41] = 1
        if self.energyTanks[6]:
            self.bitfield[42] = 1
        if self.missileTanks[19]:
            self.bitfield[43] = 1
        if self.doors["Ridley"][1]:
            self.bitfield[44] = 1
        if self.energyTanks[7]:
            self.bitfield[45] = 1
        if self.missileTanks[20]:
            self.bitfield[46] = 1
        if self.doors["Ridley"][2]:
            self.bitfield[47] = 1
        if self.energyTanks[8]:
            self.bitfield[48] = 1
        if self.missileTanks[21]:
            self.bitfield[49] = 1
        if self.doors["Tourian"][1]:
            self.bitfield[50] = 1
        if self.doors["Tourian"][2]:
            self.bitfield[51] = 1
        if self.doors["Tourian"][3]:
            self.bitfield[52] = 1
        if self.zebetitesDestroyed[1]:
            self.bitfield[53] = 1
        if self.zebetitesDestroyed[2]:
            self.bitfield[54] = 1
        if self.zebetitesDestroyed[3]:
            self.bitfield[55] = 1
        if self.zebetitesDestroyed[4]:
            self.bitfield[56] = 1
        if self.zebetitesDestroyed[5]:
            self.bitfield[57] = 1
        if self.motherBrainKilled:
            self.bitfield[58] = 1

        # 59-63 unknown

        # not 64, 65, or 66 = Brinstar
        # 64 = Norfair
        # 65 and not 66 = Kraid's Lair
        # 66 and not 65 = Ridley's Lair
        # 65 and 66 = Tourian
        if self.startLocation == 1:
            self.bitfield[64] = 1
        if self.startLocation == 2 or self.startLocation == 4:
            self.bitfield[65] = 1
        if self.startLocation == 3 or self.startLocation == 4:
            self.bitfield[66] = 1

        # 67 is the reset bit, I want all passwords to be valid
        # if self.:
        #     self.bitfield[67] = 1

        # 68-70 are unknown

        if self.swimsuit:
            self.bitfield[71] = 1
        if self.samusHas["Bombs"]:
            self.bitfield[72] = 1
        if self.samusHas["High Jump Boots"]:
            self.bitfield[73] = 1
        if self.samusHas["Long Beam"]:
            self.bitfield[74] = 1
        if self.samusHas["Screw Attack"]:
            self.bitfield[75] = 1
        if self.samusHas["Maru Mari"]:
            self.bitfield[76] = 1
        if self.samusHas["Varia"]:
            self.bitfield[77] = 1
        if self.samusHas["Wave Beam"]:
            self.bitfield[78] = 1
        if self.samusHas["Ice Beam"]:
            self.bitfield[79] = 1

        # Missile count
        # +2^n from 0 to 7
        binMissiles = bin(self.missileCount).replace('0b','')[::-1]
        while len(binMissiles) < 8:
            binMissiles += "0"
        if int(binMissiles[0]) == 1:
            self.bitfield[80] = 1
        if int(binMissiles[1]) == 1:
            self.bitfield[81] = 1
        if int(binMissiles[2]) == 1:
            self.bitfield[82] = 1
        if int(binMissiles[3]) == 1:
            self.bitfield[83] = 1
        if int(binMissiles[4]) == 1:
            self.bitfield[84] = 1
        if int(binMissiles[5]) == 1:
            self.bitfield[85] = 1
        if int(binMissiles[6]) == 1:
            self.bitfield[86] = 1
        if int(binMissiles[7]) == 1:
            self.bitfield[87] = 1
        # 88-119 are game age, leaving at 0

        # 120-123 are unknown

        # I have no idea why these are at the end. I wonder if Kraid and
        # Ridley were relatively late additions to the game? (Or maybe
        # they were just implemented late in the game's development.)
        if self.ridleyKilled:
            self.bitfield[124] = 1
        if self.ridleyStatue:
            self.bitfield[125] = 1
        if self.kraidKilled:
            self.bitfield[126] = 1
        if self.kraidStatue:
            self.bitfield[127] = 1

    def generatePassword(self):
        """ Generate the password from the bitfield.
            This is a five-step process.
            1) Reverse the order of each 8-bit byte to make it little-endian.
            2) Cycle the entire bitfield 0-7 bits to the right. 
               Append the number of shifts in binary to the end - again, little-endian.
            3) Create the checksum by turning each byte into a decimal number, 
               summing them, converting the sum *back* to binary, and taking the lowest
               8 bits of that binary sum and adding it - BIG-ENDIAN - to the end.
            4) Separate the bitstream into ***6-bit*** chunks and create a decimal
               number from each chunk (0-63).
            5) Associate each decimal number with a letter from the Metroid Alphabet
               (listed at the top of __init__()).
            
            I'm not doing step 2 yet, which is equivalent to shifting 0 places
            and making the shift byte 00000000.
        """
        # not gonna do the bit-shifting yet
        bitfield = self.bitfield
        bitfield = bitfield + [0,0,0,0,0,0,0,0] # add the zero shift byte
        self.fullbitfield = "".join([str(x) for x in bitfield])
        newBitfield = []
        for i in range(17):
            j = i * 8
            k = j + 8
            word = self.fullbitfield[j:k][::-1] # I thought [j:k:-1] should work but it doesn't
            newBitfield.append(word)
        decChecksum = sum([int(x, 2) for x in newBitfield])
        bitfield = "".join(newBitfield)
        binChecksum = bin(decChecksum).replace('0b','')
        checksum = binChecksum[-8:]
        while len(checksum) < 8:
            checksum = "0" + checksum
        for bit in checksum:
            bitfield += bit

        letters = []
        letter = []
        for bit in bitfield:
            letter.append(bit)
            if len(letter) == 6:
                letters.append(self.alphabet[int("".join(letter),2)])
                letter = []
        words = []
        word = []
        for lt in letters:
            word.append(lt)
            if len(word) == 6:
                words.append("".join(word))
                word = []
        words.append("".join(word))
        self.password = " ".join(words)
        return self.password

    def decodePassword(self, pwd):
        """ Sanity checking! This function decodes an input password back into a bitfield,
            so that you can check that it was properly encoded.
            Q: Why doesn't this display the game state?
            A: I trust that https://www.truepeacein.space properly encodes the game state.
               So when I create a game state with the randomizer, I can recreate that
               game state at TPIS and use the password generates as its input, to check
               against my randomized game password. In other words, this is a testing 
               function, and in the intended use case I'll know what the input bitfield is 
               and be able to check against it.
        """
        densePwd = pwd.replace(" ","")
        numPwd = []
        for chr in densePwd:
            numPwd.append(self.alphabet.index(chr))
        bitPwd = [bin(x).replace("0b","") for x in numPwd]
        longBitPwd = []
        for word in bitPwd:
            longword = word
            while len(longword) < 6:
                longword = "0" + longword
            longBitPwd.append(longword)
        newBitfield = "".join(longBitPwd)
        csm = sum([int(x) for x in newBitfield[:136]])
        print(csm)
        for i in range(len(newBitfield)):
            print(newBitfield[i], end="")
            if i%8 == 7:
                print(" ", end="")
            if i%64 == 63:
                print()

def main():
    gs = MetroidState()
    gs.randomize()
    gs.createBitfield()
    gs.generatePassword()
    print(gs.toString())

if __name__ == "__main__":
    if len(sys.argv) == 2:
        gs = MetroidState()
        gs.decodePassword(sys.argv[1])
    else:
        main()