utils.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wafel\n",
    "import numpy as np\n",
    "import copy\n",
    "import itertools\n",
    "from random import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VERSION = 'jp'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_inputs(g, inputs):\n",
    "    '''Set inputs for a single frame of in-game time\n",
    "    g is a Wafel game object\n",
    "    inputs is a Wafel input object'''\n",
    "    g.write('gControllerPads[0].button', inputs.buttons)\n",
    "    g.write('gControllerPads[0].stick_x', inputs.stick_x)\n",
    "    g.write('gControllerPads[0].stick_y', inputs.stick_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_actual_coord(wafel_coord):\n",
    "    '''Wafel's input specification has changed across versions. This performs the most up-to-date\n",
    "    transformation of joystick values from the standard m64 format input specification to the Wafel format.'''\n",
    "    return wafel_coord - 256*(wafel_coord > 127)\n",
    "\n",
    "def increment_coord(coord):\n",
    "    '''Helper function to ignore deadzones while incrementing a coordinate.\n",
    "    To decrement, can say x = -increment_coord(-x)'''\n",
    "    coord = to_actual_coord(coord)\n",
    "    if coord == -8:\n",
    "        coord = 0\n",
    "    elif coord == 0:\n",
    "        coord = 8\n",
    "    elif coord >= 127:\n",
    "        coord = 127\n",
    "    else:\n",
    "        coord = coord + 1\n",
    "    return coord % 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_globtime_syncs_mod64(filename, start_frame, end_frame,\n",
    "                               vars_to_sync_at_end):\n",
    "    '''For an m64 at location filename, checks whether the TAS syncs at end_frame\n",
    "    with respect to vars_to_sync_at_end if global timer is changed mod 64 at start_frame.'''\n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    for f in range(start_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    start_state = game.save_state()\n",
    "    start_globtime = game.read('gGlobalTimer')\n",
    "    for f in range(start_frame, end_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    var_values = {}\n",
    "    for var in vars_to_sync_at_end:\n",
    "        var_values[var] = game.read(var)\n",
    "    successes = []\n",
    "    for globtime_offset in range(64):\n",
    "        game.load_state(start_state)\n",
    "        game.write('gGlobalTimer', start_globtime + globtime_offset)\n",
    "        for f in range(start_frame, end_frame):\n",
    "            set_inputs(game, m64[1][f])\n",
    "            game.advance()\n",
    "        succeeded = True\n",
    "        for var in vars_to_sync_at_end:\n",
    "            if var_values[var] != game.read(var):\n",
    "                succeeded = False\n",
    "        if succeeded:\n",
    "            print(globtime_offset)\n",
    "            successes.append(globtime_offset)\n",
    "    return successes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def detect_a_presses(filename):\n",
    "    '''Counts A presses'''\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    for f in range(len(m64[1])):\n",
    "        if m64[1][f].buttons >= 0x8000:\n",
    "            print(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def identify_dustable_frames(filename, start_frame, end_frame):\n",
    "    '''For all frames with a walking action and low speed, determine\n",
    "whether it's possible to change whether Mario dusts without\n",
    "changing his position on the next frame. This is done for the sake of RNG manip.\n",
    "Exhaustively checks all joystick (x,y) pairs. This is quite slow, so I don't\n",
    "typically use it. Could presumably be rewritten to run without advancing Wafel,\n",
    "but it would require importing the trig table and writing some game logic.'''\n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    for f in range(start_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    rng_changes = []\n",
    "    for f in range(start_frame, end_frame):\n",
    "        prev_state = game.save_state()\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "        if game.read('gMarioState.action') == 67109952 and game.read('gMarioState.forwardVel') < 17:\n",
    "            rng_reached = game.read('gRandomSeed16')\n",
    "            x_reached = game.read('gMarioState.pos[0]')\n",
    "            z_reached = game.read('gMarioState.pos[2]')\n",
    "            found_soln = None\n",
    "            for x in range(256):\n",
    "                for y in range(256):\n",
    "                    game.load_state(prev_state)\n",
    "                    set_inputs(game, wafel.Input(buttons = m64[1][f].buttons, stick_x = x, stick_y = y))\n",
    "                    game.advance()\n",
    "                    if (game.read('gRandomSeed16') != rng_reached\n",
    "                        and game.read('gMarioState.pos[0]') == x_reached\n",
    "                        and game.read('gMarioState.pos[2]') == z_reached):\n",
    "                        found_soln = (x, y)\n",
    "            game.load_state(prev_state)\n",
    "            set_inputs(game, m64[1][f])\n",
    "            game.advance()\n",
    "            if found_soln != None:\n",
    "                print('Found an alternate dust ' + str(f))\n",
    "                rng_changes.append((f, m64[1][f].stick_x, m64[1][f].stick_y, found_soln[0], found_soln[1]))\n",
    "            else:\n",
    "                print('Found no alternate dust for ' + str(f))\n",
    "    return rng_changes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def identify_dustable_frames_faster(filename, start_frame, end_frame):\n",
    "    '''For all frames with a walking action and low speed, determine\n",
    "whether it's possible to change whether Mario dusts without\n",
    "changing his position on the next frame. For the sake of RNG manip.\n",
    "Only checks a subset of (x,y) pairs that should give similar angles.\n",
    "Uses actual trig instead of in-game trig table, so I can't guarantee\n",
    "that it will find all possible dustable frames. There's a window that I widened\n",
    "to a few AU to try to improve the chances that it captures all of them.\n",
    "Of course, this slows it down, but it's still much faster than the exhaustive one'''\n",
    "    sticks = np.meshgrid(np.arange(0, 256, 1), np.arange(0, 256, 1))\n",
    "    sticks = np.array([sticks[1], sticks[0]])\n",
    "    raw_sticks = copy.deepcopy(sticks)\n",
    "    #convert raw joystick inputs to processed stick inputs\n",
    "    sticks[sticks > 127] -= 256\n",
    "    for coord in range(2):\n",
    "        sticks[coord][np.abs(sticks[coord]) < 8] = 0\n",
    "        sticks[coord][sticks[coord] >= 8] -= 6\n",
    "        sticks[coord][sticks[coord] <= -8] += 6\n",
    "    angles = np.arctan2(sticks[0], sticks[1])\n",
    "    \n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    for f in range(start_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    rng_changes = []\n",
    "    for f in range(start_frame, end_frame):\n",
    "        prev_state = game.save_state()\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "        if game.read('gMarioState.action') == 67109952 and game.read('gMarioState.forwardVel') < 17:\n",
    "            angle = angles[m64[1][f].stick_x, m64[1][f].stick_y]\n",
    "            \n",
    "            \n",
    "            rng_reached = game.read('gRandomSeed16')\n",
    "            x_reached = game.read('gMarioState.pos[0]')\n",
    "            z_reached = game.read('gMarioState.pos[2]')\n",
    "            found_soln = None\n",
    "            \n",
    "            \n",
    "            #check for positions with angle within about 2 HAU (extra slack just in case due to\n",
    "            #not using sm64 trig tables)\n",
    "            similar_angle_joys = np.where(np.abs(angles - angle) < .0091)\n",
    "            for attempt in range(similar_angle_joys[0].shape[0]):\n",
    "                game.load_state(prev_state)\n",
    "                try_x = raw_sticks[0][similar_angle_joys[0][attempt], similar_angle_joys[1][attempt]]\n",
    "                try_y = raw_sticks[1][similar_angle_joys[0][attempt], similar_angle_joys[1][attempt]]\n",
    "\n",
    "                set_inputs(game, wafel.Input(buttons = m64[1][f].buttons, stick_x = try_x, stick_y = try_y))\n",
    "                game.advance()\n",
    "                if (game.read('gRandomSeed16') != rng_reached\n",
    "                    and game.read('gMarioState.pos[0]') == x_reached\n",
    "                    and game.read('gMarioState.pos[2]') == z_reached):\n",
    "                    found_soln = (try_x, try_y)\n",
    "                    break            \n",
    "            \n",
    "            game.load_state(prev_state)\n",
    "            set_inputs(game, m64[1][f])\n",
    "            game.advance()\n",
    "            if found_soln != None:\n",
    "                print('Found an alternate dust ' + str(f))\n",
    "                rng_changes.append((f, m64[1][f].stick_x, m64[1][f].stick_y, found_soln[0], found_soln[1]))\n",
    "            else:\n",
    "                print('Found no alternate dust for ' + str(f))\n",
    "    return rng_changes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rng_advance():\n",
    "    '''rng behavior based on Ukikipedia pseudocode'''\n",
    "    global rngseed\n",
    "    if rngseed == 22026:\n",
    "        rngseed = np.uint16(0)\n",
    "    \n",
    "    A = np.uint16(rngseed << 8)\n",
    "    A = A ^ rngseed\n",
    "    \n",
    "    rngseed = A\n",
    "    left_byte = np.uint16(A << 8)\n",
    "    right_byte = A >> 8\n",
    "    rngseed = left_byte + right_byte\n",
    "    \n",
    "    A = np.uint16(A << 8) >> 7\n",
    "    A = A ^ rngseed\n",
    "    B = A >> 1\n",
    "    B = B ^ np.uint16(65408)\n",
    "    \n",
    "    if A % 2 == 0:\n",
    "        if B == 43605:\n",
    "            rngseed = 0\n",
    "        else:\n",
    "            rngseed = B ^ np.uint16(8180)\n",
    "    else:\n",
    "        rngseed = B ^ np.uint16(33152)\n",
    "    return rngseed\n",
    "\n",
    "def get_rng_array():\n",
    "    '''Computes the array of RNG values used by sm64. We tend to\n",
    "    use the indices while communicating about RNG rather than the values.\n",
    "    They are more useful as they give a sense of proximity to the RNG you are targeting.'''\n",
    "    global rngseed\n",
    "    rngseed = np.uint16(0)\n",
    "    values = []\n",
    "    while rngseed > 0 or len(values) == 0:\n",
    "        values.append(rngseed)\n",
    "        rng_advance()\n",
    "    return values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_rng_syncs(filename, start_frame, end_frame,\n",
    "                               vars_to_sync_at_end,\n",
    "                    low = 0, high = 65114):\n",
    "    '''For a TAS at filename, tries setting each RNG at start_frame and then advances\n",
    "    to end_frame. Checks whether values in vars_to_sync_at_end (a dictionary from variables\n",
    "    to values) align to the original TAS.'''\n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    for f in range(start_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    start_state = game.save_state()\n",
    "    for f in range(start_frame, end_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    var_values = {}\n",
    "    for var in vars_to_sync_at_end:\n",
    "        var_values[var] = game.read(var)\n",
    "    successes = []\n",
    "    game.write('gRandomSeed16', 0)\n",
    "    print(low)\n",
    "    print(high)\n",
    "    for ind, val in enumerate(get_rng_array()[low:high]):\n",
    "        game.load_state(start_state)\n",
    "        game.write('gRandomSeed16', int(val))\n",
    "        for f in range(start_frame, end_frame):\n",
    "            set_inputs(game, m64[1][f])\n",
    "            game.advance()\n",
    "        succeeded = True\n",
    "        for var in vars_to_sync_at_end:\n",
    "            if var_values[var] != game.read(var):\n",
    "                succeeded = False\n",
    "        if succeeded:\n",
    "            print(ind+low)\n",
    "            successes.append(ind+low)\n",
    "    return successes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_fitness(game, m64, start_state,\n",
    "                     start_frame, end_frame,\n",
    "                     possible_alternatives,\n",
    "                     goal_rng_ind, possible_alternatives_choices,\n",
    "                     desired_x, desired_z,\n",
    "                     rng_val_to_ind_map):\n",
    "    '''For internal use in the target_rng function.\n",
    "    Essentially, checks how close a sequences of choices to dust or not dust\n",
    "    is to obtaining the desired RNG index while still keeping Mario at his desired position.\n",
    "    The fitness is essentially just distance between RNG indices. However, there is a penalty\n",
    "    (currently 32) for being off mod 4 as it is much more difficult to change RNG mod 4.'''\n",
    "    game.load_state(start_state)\n",
    "    cur_alternative_ind = 0\n",
    "    for f in range(start_frame, end_frame):\n",
    "        if cur_alternative_ind < len(possible_alternatives) and f == possible_alternatives[cur_alternative_ind][0]:\n",
    "            if possible_alternatives_choices[cur_alternative_ind] == 1:\n",
    "                x, y = possible_alternatives[cur_alternative_ind][3:]\n",
    "                set_inputs(game, wafel.Input(m64[1][f].buttons, x, y))\n",
    "            else:\n",
    "                set_inputs(game, m64[1][f])\n",
    "            cur_alternative_ind += 1\n",
    "        else:\n",
    "            set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    rng_result = rng_val_to_ind_map[game.read('gRandomSeed16')]\n",
    "    cost = abs(rng_result - goal_rng_ind)\n",
    "    if cost > 33000:\n",
    "        cost = abs(cost - 65114)\n",
    "    if desired_x != game.read('gMarioState.pos[0]') or desired_z != game.read('gMarioState.pos[2]'):\n",
    "        cost += 10000000\n",
    "    if (rng_result - goal_rng_ind) % 4 != 0:\n",
    "        # Penalty for being off mod 4\n",
    "        cost += 32\n",
    "    return cost\n",
    "    \n",
    "        \n",
    "\n",
    "\n",
    "def target_rng(filename, filename_output, start_frame, end_frame, goal_rng_ind,\n",
    "               max_attempts = 100000, starting_temp = 3, change_prob = .3):\n",
    "    '''Tries to manipulate dust from start_frame to end_frame to achieve a\n",
    "    certain RNG index at end_frame. Checks Mario's position to verify it didn't\n",
    "    desync.\n",
    "    filename is the input m64, and a solution if found will be written to filename_output\n",
    "    goal_rng_ind is the desired index\n",
    "    The optimization is done via annealing in the 2^n space (where n is the number of frames\n",
    "    where dust manip is possible). starting_temp is the temperature used for this annealing.\n",
    "    change_prob determines with what probability each dust-no-dust possibility in an interval\n",
    "    is flipped for each attempt.\n",
    "    max_attempts is the number of dust-no-dust arrays it will attempt to check before giving up'''\n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    for f in range(start_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    start_state = game.save_state()\n",
    "    \n",
    "    # RNG array: mapping from indices to values. Interested in the reverse\n",
    "    rng_array = get_rng_array()\n",
    "    rng_val_to_ind_map = {rng_array[ind]: ind for ind in range(65114)}\n",
    "    \n",
    "    # Run the game and save desired Mario position at the end\n",
    "    for f in range(start_frame, end_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "        \n",
    "    desired_x = game.read('gMarioState.pos[0]')\n",
    "    desired_z = game.read('gMarioState.pos[2]')\n",
    "    game.load_state(start_state)\n",
    "    \n",
    "    # Determine on which frames we can RNG manip\n",
    "    possible_alternatives = identify_dustable_frames_faster(filename, start_frame, end_frame)\n",
    "    print('Identified {0} frames on which one can do RNG manip'.format(len(possible_alternatives)))\n",
    "    \n",
    "    cur_choices = np.zeros(len(possible_alternatives))\n",
    "    cur_fitness = evaluate_fitness(game, m64, start_state,\n",
    "                                    start_frame, end_frame,\n",
    "                                    possible_alternatives,\n",
    "                                    goal_rng_ind, cur_choices,\n",
    "                                    desired_x, desired_z,\n",
    "                                    rng_val_to_ind_map)\n",
    "\n",
    "    # Annealing\n",
    "    temp = starting_temp\n",
    "    soln = None\n",
    "    for attempts in range(max_attempts):\n",
    "        if attempts % 100 == 0:\n",
    "            print(attempts)\n",
    "            print(cur_fitness)\n",
    "            temp *= .98\n",
    "            if change_prob*len(possible_alternatives) > 3:\n",
    "                change_prob *= .98\n",
    "        to_change = np.random.choice(a=[0, 1], size=(len(possible_alternatives)), p=[1 - change_prob, change_prob])\n",
    "        # Sometimes, we may only want to change close to the end, as making changes near to the end is less\n",
    "        # likely to lead to a desync. With some probability, we do this. The length of the postfix is\n",
    "        # also chosen randomly.\n",
    "        if random() < .3:\n",
    "            to_change[0:int(random()*len(possible_alternatives))] = 0\n",
    "        #Apply the changes and evaluate the fitness\n",
    "        new_choices = (cur_choices + to_change) % 2\n",
    "        new_fitness = evaluate_fitness(game, m64, start_state,\n",
    "                                       start_frame, end_frame,\n",
    "                                       possible_alternatives,\n",
    "                                       goal_rng_ind, new_choices,\n",
    "                                       desired_x, desired_z,\n",
    "                                       rng_val_to_ind_map)\n",
    "        if new_fitness == 0:\n",
    "            soln = new_choices\n",
    "            print('Found an RNG resync solution.')\n",
    "            print(game.read('gRandomSeed16'))\n",
    "            break\n",
    "        if new_fitness < cur_fitness or random() < np.exp(-(new_fitness - cur_fitness) / temp):\n",
    "            cur_choices = new_choices\n",
    "            cur_fitness = new_fitness\n",
    "    if soln is None:\n",
    "        print('Could not find a solution. Final fitness is {0}'.format(cur_fitness))\n",
    "        soln = cur_choices\n",
    "    for choice_ind in range(len(possible_alternatives)):\n",
    "        if soln[choice_ind] == 1:\n",
    "            alt = possible_alternatives[choice_ind]\n",
    "            m64[1][alt[0]] = wafel.Input(m64[1][alt[0]].buttons, alt[3], alt[4])\n",
    "    \n",
    "    wafel.save_m64(filename_output, m64[0], m64[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_working_indices_with_manip(filename, start_frame, end_frame, check_vars = [], custom_checkpoints_to_add = [],\n",
    "                                    inds_to_check = range(65114), consecutive_reduction = False, num_files_to_save = 10,\n",
    "                                    hack_index = None, checkpoints_to_omit = []):\n",
    "    '''Does a search of possible RNG manipulations starting at start_frame to try to match the given variables\n",
    "    at end_frame for different starting RNG indices. Stores the first num_files_to_save files that it manages to get working.\n",
    "    Checks Mario's position and variables in check_vars to determine whether the TAS has synced.\n",
    "    Checks the variables every 1/10 of the interval from start_frame to end_frame. You can add your own checkpoints\n",
    "    via custom_checkpoints_to_add or suppress them by supplying an array in checkpoints_to_omit. The checkpoints help\n",
    "    speed up the execution via early termination.\n",
    "    inds_to_check is just the set of indices for which you want to search for syncing TASes.\n",
    "    consecutive_reduction speeds things up at the cost of thoroughness. If RNG manip is possible on 6 consecutive\n",
    "    frames, consecutive_reduction=False will check all 64 possibilities, but True will check only 7 (change no frames,\n",
    "    change only first, change only first two...)\n",
    "    hack_index hacks the RNG index at start_frame in case the TAS doesn't actually sync for the RNG you currently have,\n",
    "    but you know one for which it does\n",
    "    if use_checkpoints is set to false, doesn't use checkpoints at all'''\n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    for f in range(start_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    \n",
    "    rng_array = get_rng_array()\n",
    "    if hack_index != None:\n",
    "        game.write('gRandomSeed16', int(rng_array[hack_index]))\n",
    "    start_state = game.save_state()\n",
    "    \n",
    "    num_files_saved = 0\n",
    "    \n",
    "    # Find appropriate values at checkpoints\n",
    "    check_points = np.linspace(start_frame, end_frame - 1, 10)\n",
    "    check_points = set([int(f) for f in check_points] + custom_checkpoints_to_add)\n",
    "    check_points = set([f for f in check_points if f not in checkpoints_to_omit])\n",
    "    if not use_checkpoints:\n",
    "        check_points = set([])\n",
    "    correct_vals_map = dict()\n",
    "    for f in range(start_frame, end_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        if f in check_points:\n",
    "            correct_vals_map[f] = dict()\n",
    "            for var in check_vars:\n",
    "                correct_vals_map[f][var] = game.read(var)\n",
    "        game.advance()\n",
    "    print(correct_vals_map)\n",
    "        \n",
    "    desired_x = game.read('gMarioState.pos[0]')\n",
    "    desired_z = game.read('gMarioState.pos[2]')\n",
    "    game.load_state(start_state)\n",
    "    \n",
    "    # Determine on which frames we can manip\n",
    "    possible_alternatives = identify_dustable_frames_faster(filename, start_frame, end_frame)\n",
    "    print('Identified {0} frames on which one can do RNG manip'.format(len(possible_alternatives)))\n",
    "    \n",
    "    # we might have situations where there are like 6 consecutive frames to manip\n",
    "    # instead of checking all 64, we can just try yes*k and then no*(6-k) for all k\n",
    "    # Should be fine to take itertools.product and then filter them out after, since frame\n",
    "    # advancing will be a much slower step\n",
    "    # Store the possibilities\n",
    "    to_check_possibilities = []\n",
    "    if consecutive_reduction:\n",
    "        for possible_alternatives_choices in itertools.product([0, 1], repeat = len(possible_alternatives)):\n",
    "            skip = False\n",
    "            for i in range(len(possible_alternatives_choices)):\n",
    "                if i > 0 and (possible_alternatives_choices[i] == 1 and\n",
    "                              possible_alternatives_choices[i - 1] == 0 and\n",
    "                              possible_alternatives[i][0] == possible_alternatives[i - 1][0] + 1):\n",
    "                    skip = True\n",
    "                    break\n",
    "            if not skip:\n",
    "                to_check_possibilities.append(possible_alternatives_choices)\n",
    "    \n",
    "    working = []\n",
    "    \n",
    "    # Now check whether each index can be made to work\n",
    "    for count, rng_ind in enumerate(inds_to_check):\n",
    "        # Maintain a list of failed prefixes (that is, situations where we fail at a checkpoint only\n",
    "        # after some of the dust manip opportunities) to prune more possibilities\n",
    "        failed_prefixes = set()\n",
    "        choices_iterator = itertools.product([0, 1], repeat = len(possible_alternatives))\n",
    "        if consecutive_reduction:\n",
    "            choices_iterator = to_check_possibilities\n",
    "        # Now just run through each option and check whether it worked\n",
    "        for possible_alternatives_choices in choices_iterator:        \n",
    "            if count % 100 == 0 and max(possible_alternatives_choices) == 0:\n",
    "                print(count)\n",
    "            skip = False\n",
    "            for i in range(len(possible_alternatives_choices)):\n",
    "                # Pruning step if this prefix is already known to be a failure\n",
    "                if tuple(possible_alternatives_choices[:i]) in failed_prefixes:\n",
    "                    skip = True\n",
    "                    break\n",
    "            if skip:\n",
    "                continue\n",
    "            game.load_state(start_state)\n",
    "            game.write('gRandomSeed16', int(rng_array[rng_ind]))\n",
    "            \n",
    "            cur_alternative_ind = 0\n",
    "            succeeded = True\n",
    "            # Run the game for this set of possibilities\n",
    "            for f in range(start_frame, end_frame):\n",
    "                # Check for failure if we are at a checkpoint\n",
    "                if f in check_points:\n",
    "                    for var in check_vars:\n",
    "                        if correct_vals_map[f][var] != game.read(var):\n",
    "                            failed_prefixes.add(tuple(possible_alternatives_choices[:cur_alternative_ind]))\n",
    "                            succeeded = False\n",
    "                            break\n",
    "                if not succeeded:\n",
    "                    break\n",
    "                # Decide what joystick input to have if we are at a manipable frame\n",
    "                if cur_alternative_ind < len(possible_alternatives) and f == possible_alternatives[cur_alternative_ind][0]:\n",
    "                    if possible_alternatives_choices[cur_alternative_ind] == 1:\n",
    "                        x, y = possible_alternatives[cur_alternative_ind][3:]\n",
    "                        set_inputs(game, wafel.Input(m64[1][f].buttons, x, y))\n",
    "                    else:\n",
    "                        set_inputs(game, m64[1][f])\n",
    "                    cur_alternative_ind += 1\n",
    "                else: # Otherwise, just do the actual input in the m64\n",
    "                    set_inputs(game, m64[1][f])\n",
    "                game.advance()\n",
    "            # Finished running the file. If we're at the right position, we have a solution\n",
    "            if desired_x != game.read('gMarioState.pos[0]') or desired_z != game.read('gMarioState.pos[2]'):\n",
    "                succeeded = False\n",
    "            if succeeded:\n",
    "                print(str(possible_alternatives) + ' ' + str(rng_ind))\n",
    "                working.append(rng_ind)\n",
    "                # Produce the resulting m64 and save it\n",
    "                if num_files_saved < num_files_to_save:\n",
    "                    num_files_saved += 1\n",
    "                    m64_modify = wafel.load_m64(filename)\n",
    "                    for choice_ind in range(len(possible_alternatives)):\n",
    "                        if possible_alternatives_choices[choice_ind] == 1:\n",
    "                            alt = possible_alternatives[choice_ind]\n",
    "                            m64_modify[1][alt[0]] = wafel.Input(m64_modify[1][alt[0]].buttons, alt[3], alt[4])\n",
    "                    wafel.save_m64('generated_m64s/hack_index_{0}_frame_{1}.m64'.format(rng_ind, start_frame + 1), m64_modify[0], m64_modify[1])\n",
    "                break\n",
    "                \n",
    "    return working"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inspect(filename, frame_number_start, frame_number_end, extra_vars = []):\n",
    "    '''Inspects memory values at each frame number from frame_number_start to frame_number_end\n",
    "    for the TAS at filename. Has some default variables to check and you can add your own in extra_vars\n",
    "    Just prints it all out'''\n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    for f in range(frame_number_start):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    for f in range(frame_number_start, frame_number_end):\n",
    "        print('Frame number: {0}'.format(f))\n",
    "        print('Mario position: {0}'.format(game.read('gMarioState.pos')))\n",
    "        print('Mario yaw: {0}'.format(game.read('gMarioState.faceAngle[1]')))\n",
    "        print('Mario hspd: {0}'.format(game.read('gMarioState.forwardVel')))\n",
    "        for extra_var in extra_vars:\n",
    "            print('{0}: {1}'.format(extra_var, game.read(extra_var)))\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def locate_framenums_close(filename, mario_x, mario_y, mario_z, radius = 100):\n",
    "    '''For the TAS at filename, prints the frame number and Mario's position at every\n",
    "    frame for which his position is within 100 units of (mario_x, mario_y, mario_z)\n",
    "    Can be helpful when you want to find out on which frame Mario enters a level or something\n",
    "    without playing the whole m64 back'''\n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    for f in range(len(m64[1])):\n",
    "        mario_pos = game.read('gMarioState.pos')\n",
    "        dist = ((mario_pos[0] - mario_x)**2 + (mario_pos[1] - mario_y)**2 + (mario_pos[2] - mario_z)**2)**.5\n",
    "        if dist < radius:\n",
    "            print('{0}: {1} pos, {2} stars'.format(f, mario_pos, game.read('gMarioState.numStars')))\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#******************************\n",
    "#NOT AT ALL READY YET. Have not worked on this in a while, may never finish it\n",
    "#Idea was to try to take two TASes where the final star collected by one is faster\n",
    "#than the same star in the other, replace the inputs, and try to sync RNG after.\n",
    "#This may often fail to fix things though (due to not being able to access the desired RNG or due to\n",
    "#non-RNG issues like global timer/twirl yaw/HOLP/etc), and doesn't consider things like textboxes, so it\n",
    "#may not be worth trying to finish since it will so frequently not work.\n",
    "def hex_in_improvement(filename_old, filename_improved, filename_output,\n",
    "                       frames_back_to_search_slowwalks = 5000):\n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    start_state = game.save_state()\n",
    "    m64_old = wafel.load_m64(filename_old)\n",
    "    m64_improved = wafel.load_m64(filename_improved)\n",
    "    print(len(m64_old[1]))\n",
    "    print(len(m64_improved[1]))\n",
    "    \n",
    "    rng_array = get_rng_array()\n",
    "    rng_array_inv = {rng_array[i]: i for i in range(len(rng_array))}\n",
    "    \n",
    "    #collect relevant information about the TASes\n",
    "    #First, identify which frame we should sync on by finding last frame before Mario stops\n",
    "    #moving. Measure position and number of stars on this frame; the new TAS should match this.\n",
    "    prev_mario_pos, goal_mario_pos, num_stars, frame_to_sync_on_improved, actual_index = None, None, None, None, None\n",
    "    for f in range(len(m64_improved[1])):\n",
    "        set_inputs(game, m64_improved[1][f])\n",
    "        game.advance()\n",
    "        new_mario_pos = game.read('gMarioState.pos')\n",
    "        if new_mario_pos != prev_mario_pos:\n",
    "            num_stars = game.read('gMarioState.numStars')\n",
    "            frame_to_sync_on_improved = f\n",
    "            goal_mario_pos = new_mario_pos\n",
    "            actual_index = rng_array_inv[game.read('gRandomSeed16')]\n",
    "        prev_mario_pos = new_mario_pos\n",
    "        \n",
    "    print(num_stars)\n",
    "       \n",
    "    \n",
    "    \n",
    "    game.load_state(start_state)\n",
    "    frame_to_sync_on_old, orig_rng_index = None, None\n",
    "    for f in range(len(m64_old[1])):\n",
    "        set_inputs(game, m64_old[1][f])\n",
    "        game.advance()\n",
    "        if game.read('gMarioState.pos') == goal_mario_pos and game.read('gMarioState.numStars') == num_stars:\n",
    "            frame_to_sync_on_old = f\n",
    "            orig_rng_index = rng_array_inv[game.read('gRandomSeed16')]\n",
    "            break\n",
    "    \n",
    "    print('Try to synchronize after {0} stars with pos {1} on frame {2}'.format(num_stars, goal_mario_pos, frame_to_sync_on_improved))\n",
    "            \n",
    "    print('Want to reach rng index {0}'.format(orig_rng_index))\n",
    "    print('Actually observed index: {0}'.format(actual_index))\n",
    "    hexed_inputs_wrongrng = m64_improved[1][:frame_to_sync_on_improved] + m64_old[1][frame_to_sync_on_old:]\n",
    "    filename_intermediate = filename_output[:-4] + '_with_ending_hexed_badrng.m64'\n",
    "    wafel.save_m64(filename_intermediate, m64_improved[0], hexed_inputs_wrongrng)\n",
    "    \n",
    "    target_rng(filename_intermediate, filename_output,\n",
    "               frame_to_sync_on_improved - frames_back_to_search_slowwalks, frame_to_sync_on_improved,\n",
    "               orig_rng_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def aim_for_target(filename, start_frame, end_frame, vars_list, var_targets,\n",
    "                       low = 0, high = 65114, repair_mario_positions = False,\n",
    "                   alternate_repair_source = None):\n",
    "    '''Try all possible rng_indices at start_frame and at end_frame evaluate\n",
    "       how close the variables are to the targets.\n",
    "       If Mario's positions might desync as a result of the RNG, force him to be back in the\n",
    "       same place so that he (probably) influences enemies the same way\n",
    "       If repair_source is None, will repair Mario's position to what is actually\n",
    "       observed in the TAS originally.\n",
    "       Otherwise, alternate_repair_source should be a tuple (filename, start_frame).\n",
    "       low and high are just the minimum and maximum indices to be checked\n",
    "       Uses taxicab dist'''\n",
    "        \n",
    "    game = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "    m64 = wafel.load_m64(filename)\n",
    "    mario_positions = []\n",
    "    for f in range(start_frame):\n",
    "        set_inputs(game, m64[1][f])\n",
    "        game.advance()\n",
    "    start_state = game.save_state()\n",
    "    if repair_mario_positions:\n",
    "        if alternate_repair_source is None:\n",
    "            for f in range(start_frame, end_frame):\n",
    "                set_inputs(game, m64[1][f])\n",
    "                game.advance()\n",
    "                mario_positions.append(game.read('gMarioState.pos'))\n",
    "        else:\n",
    "            game_alt = wafel.Game('sm64_{0}.dll'.format(VERSION))\n",
    "            m64_alt = wafel.load_m64(alternate_repair_source[0])\n",
    "            for f in range(alternate_repair_source[1]):\n",
    "                set_inputs(game_alt, m64_alt[1][f])\n",
    "                game_alt.advance()\n",
    "            for f in range(alternate_repair_source[1], alternate_repair_source[1] + end_frame - start_frame):\n",
    "                set_inputs(game_alt, m64_alt[1][f])\n",
    "                game_alt.advance()\n",
    "                mario_positions.append(game_alt.read('gMarioState.pos'))\n",
    "    #print(mario_positions)\n",
    "    distances = []\n",
    "    rng_array = get_rng_array()\n",
    "    for rng_ind in range(low, high):\n",
    "        rng_val = int(rng_array[rng_ind])\n",
    "        if rng_ind % 1000 == 0 and rng_ind > low:\n",
    "            print(rng_ind)\n",
    "            print(min(distances))\n",
    "        game.load_state(start_state)\n",
    "        game.write('gRandomSeed16', rng_val)\n",
    "        for f in range(start_frame, end_frame):\n",
    "            set_inputs(game, m64[1][f])\n",
    "            game.advance()\n",
    "            if repair_mario_positions:\n",
    "                game.write('gMarioState.pos[0]', mario_positions[f - start_frame][0])\n",
    "                game.write('gMarioState.pos[1]', mario_positions[f - start_frame][1])\n",
    "                game.write('gMarioState.pos[2]', mario_positions[f - start_frame][2])\n",
    "        error = 0\n",
    "        for i in range(len(vars_list)):\n",
    "            actual = game.read(vars_list[i])\n",
    "            error += abs(actual - var_targets[i])\n",
    "        distances.append(error)\n",
    "    return distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}