package main

import (
	"log"
	"math"
	"os"
	"strings"

	"github.com/thoas/go-funk"
)

func main() {
	inputBytes, _ := os.ReadFile("input.txt")
	inputs := strings.Split(string(inputBytes), "\n")

	log.Println("Puzzle 1:", puzzle1(inputs))
	log.Println("Puzzle 2:", puzzle2(inputs))
}

func puzzle1(input []string) int {
	// Each bit in the gamma rate can be determined by finding the most common bit in the corresponding position of all numbers in the diagnostic report.
	// The epsilon rate is calculated in a similar way; rather than use the most common bit, the least common bit from each position is used.
	zeros := make([]int, len(input[0]))
	ones := make([]int, len(input[0]))
	for _, line := range input {
		for i, c := range line {
			if c == '0' {
				zeros[i]++
			} else if c == '1' {
				ones[i]++
			}
		}
	}
	// Get most and least common
	mostCommon := make([]int, len(input[0]))
	leastCommon := make([]int, len(input[0]))
	for i, z := range zeros {
		if z > ones[i] {
			mostCommon[i] = 0
			leastCommon[i] = 1
		} else {
			mostCommon[i] = 1
			leastCommon[i] = 0
		}
	}
	// Convert mostCommon and leastCommon from bits to ints
	gammaRate := 0
	for i, c := range funk.ReverseInt(mostCommon) {
		gammaRate += c * int(math.Pow(2, float64(i)))
	}
	epsilonRate := 0
	for i, c := range funk.ReverseInt(leastCommon) {
		epsilonRate += c * int(math.Pow(2, float64(i)))
	}
	// Return the product of the two rates, the power consumption
	return gammaRate * epsilonRate
}

func puzzle2(input []string) int {
	// Most common bit on position
	mostCommon := func(input []string, position int, whenEqual int) int {
		zeros := 0
		ones := 0
		for _, line := range input {
			if line[position] == '0' {
				zeros++
			} else if line[position] == '1' {
				ones++
			}
		}
		if zeros > ones {
			return 0
		} else if ones > zeros {
			return 1
		}
		return whenEqual
	}
	// Filter inputs
	inputLen := len(input[0])
	oxygen := input
	for i := 0; i < inputLen; i++ {
		if len(oxygen) < 2 {
			break
		}
		most := mostCommon(oxygen, i, 1)
		oxygen = funk.FilterString(oxygen, func(s string) bool {
			if most == 1 {
				return s[i] == '1'
			} else {
				return s[i] == '0'
			}
		})
	}
	co2 := input
	for i := 0; i < inputLen; i++ {
		if len(co2) < 2 {
			break
		}
		least := 1 - mostCommon(co2, i, 1)
		co2 = funk.FilterString(co2, func(s string) bool {
			if least == 1 {
				return s[i] == '1'
			} else {
				return s[i] == '0'
			}
		})
	}
	// Convert first string in oxygen from bits to int
	oxygenRating := 0
	for i, c := range funk.ReverseString(oxygen[0]) {
		oxygenRating += int(c-'0') * int(math.Pow(2, float64(i)))
	}
	// Convert first string in co2 from bits to int
	co2Rating := 0
	for i, c := range funk.ReverseString(co2[0]) {
		co2Rating += int(c-'0') * int(math.Pow(2, float64(i)))
	}
	// Return the product of the two rates, the life support rating
	return oxygenRating * co2Rating
}