package main

import (
	"bufio"
	"errors"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
)

func main() {
	input := stdinToStringSlice()
	pt2Idx := part1(input)
	part2(input[pt2Idx:])
}

var probs map[string]int

func part1(input []string) int {
	probs = make(map[string]int)
	var messySamples int
	var ret int
	for i := 0; i < len(input); i += 4 {
		if !strings.HasPrefix(input[i], "Before: ") {
			// We're done with the part 1 input
			ret = i
			break
		}
		aft := NewState()
		aft.loadState(input[i+2])
		if cmd, err := aft.parseCommand(input[i+1]); err == nil {
			origCd := cmd[0]
			var numMatches int
			for j := 0; j < 16; j++ {
				bef := NewState()
				bef.loadState(input[i])
				cmd[0] = j
				bef.run(cmd)
				if bef.registerString() == aft.registerString() {
					numMatches++
					key := fmt.Sprintf("%d:%d", origCd, j)
					probs[key]++
				}
			}
			if numMatches > 2 {
				messySamples++
			}
		}
	}
	fmt.Println("= Part 1 =")
	fmt.Println(messySamples)
	return ret + 2
}

var opMap [16]int

func part2(input []string) {
	for i := 0; i < 16; i++ {
		opMap[i] = -1
	}
	fmt.Println("= Part 2 =")
	m1, m2 := -1, -1
	for {
		m1, m2, probs = findOnlyChance(probs)
		if m1 == -1 {
			break
		}
		opMap[m1] = m2
	}

	vpc := NewState()
	for i := range input {
		if cmd, err := vpc.parseCommand(input[i]); err == nil {
			cmd[0] = opMap[cmd[0]]
			vpc.run(cmd)
		}
	}
	fmt.Println(vpc.reg[0])
}

func findOnlyChance(of map[string]int) (int, int, map[string]int) {
	m1, m2 := -1, -1
	for i := 0; i < 16; i++ {
		var num int
		for k := range of {
			codes := strings.Split(k, ":")
			if Atoi(codes[0]) == i {
				num++
				if num > 1 {
					m2 = -1
					break
				}
				m2 = Atoi(codes[1])
			}
		}
		if num == 1 {
			m1 = i
			break
		}
		if num > 1 {
			continue
		}
	}
	if m1 == -1 {
		return -1, -1, of
	}
	// If we're here, we found something
	// Remove other possibilities for m1: and :m2 and return
	for k := range of {
		a1 := strconv.Itoa(m1)
		a2 := strconv.Itoa(m2)
		if strings.HasPrefix(k, a1+":") ||
			strings.HasSuffix(k, ":"+a2) {
			delete(of, k)
		}
	}
	return m1, m2, of
}

type State struct {
	reg     [4]int
	OpCodes map[int]func(a, b, c int)
}

func NewState() *State {
	s := &State{}
	s.OpCodes = make(map[int]func(a, b, c int))
	s.OpCodes[0] = s.addr
	s.OpCodes[1] = s.addi
	s.OpCodes[2] = s.mulr
	s.OpCodes[3] = s.muli
	s.OpCodes[4] = s.banr
	s.OpCodes[5] = s.bani
	s.OpCodes[6] = s.borr
	s.OpCodes[7] = s.bori
	s.OpCodes[8] = s.setr
	s.OpCodes[9] = s.seti
	s.OpCodes[10] = s.gtir
	s.OpCodes[11] = s.gtri
	s.OpCodes[12] = s.gtrr
	s.OpCodes[13] = s.eqir
	s.OpCodes[14] = s.eqri
	s.OpCodes[15] = s.eqrr
	return s
}

func (s *State) registerString() string {
	return fmt.Sprint(s.reg)
}

func (s *State) loadState(inp string) {
	s.reg = [4]int{0, 0, 0, 0}
	inp = strings.Trim(inp, "ABeftor: []")
	pts := strings.Split(inp, ",")
	for i := range pts {
		s.reg[i] = Atoi(strings.TrimSpace(pts[i]))
	}
}

func (s *State) parseCommand(inp string) ([]int, error) {
	var ret []int
	pts := strings.Split(inp, " ")
	for i := range pts {
		ret = append(ret, Atoi(pts[i]))
	}
	if len(ret) != 4 {
		return ret, errors.New("Invalid Command: " + inp)
	}
	return ret, nil
}

func (s *State) run(cmd []int) error {
	if fn, ok := s.OpCodes[cmd[0]]; !ok {
		return errors.New(fmt.Sprintf("Invalid OpCode: %d", cmd[0]))
	} else {
		fn(cmd[1], cmd[2], cmd[3])
	}
	return nil
}

func (s *State) addr(r1, r2, r3 int) {
	s.reg[r3] = s.reg[r1] + s.reg[r2]
}

func (s *State) addi(r1, v2, r3 int) {
	s.reg[r3] = s.reg[r1] + v2
}

func (s *State) mulr(r1, r2, r3 int) {
	s.reg[r3] = s.reg[r1] * s.reg[r2]
}

func (s *State) muli(r1, v2, r3 int) {
	s.reg[r3] = s.reg[r1] * v2
}

func (s *State) banr(r1, r2, r3 int) {
	s.reg[r3] = s.reg[r1] & s.reg[r2]
}

func (s *State) bani(r1, v2, r3 int) {
	s.reg[r3] = s.reg[r1] & v2
}

func (s *State) borr(r1, r2, r3 int) {
	s.reg[r3] = s.reg[r1] | s.reg[r2]
}

func (s *State) bori(r1, v2, r3 int) {
	s.reg[r3] = s.reg[r1] | v2
}

func (s *State) setr(r1, r2, r3 int) {
	s.reg[r3] = s.reg[r1]
}

func (s *State) seti(v1, v2, r3 int) {
	s.reg[r3] = v1
}

func (s *State) gtir(v1, r2, r3 int) {
	s.reg[r3] = boolToInt(v1 > s.reg[r2])
}

func (s *State) gtri(r1, v2, r3 int) {
	s.reg[r3] = boolToInt(s.reg[r1] > v2)
}

func (s *State) gtrr(r1, r2, r3 int) {
	s.reg[r3] = boolToInt(s.reg[r1] > s.reg[r2])
}

func (s *State) eqir(v1, r2, r3 int) {
	s.reg[r3] = boolToInt(v1 == s.reg[r2])
}

func (s *State) eqri(r1, v2, r3 int) {
	s.reg[r3] = boolToInt(s.reg[r1] == v2)
}

func (s *State) eqrr(r1, r2, r3 int) {
	s.reg[r3] = boolToInt(s.reg[r1] == s.reg[r2])
}

func boolToInt(v bool) int {
	if v {
		return 1
	}
	return 0
}

func stdinToStringSlice() []string {
	var input []string
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		input = append(input, scanner.Text())
	}
	return input
}

func Atoi(i string) int {
	var ret int
	var err error
	if ret, err = strconv.Atoi(i); err != nil {
		log.Fatal("Invalid Atoi: " + i)
	}
	return ret
}