Made a few changes, states still generating incorrectly, however input from states 0-4 looks close