// compute probability of landing on tile after certain number of steps of brownian
package main

import (
	"flag"
	"fmt"
	"math/big"
	"math/rand"
	"os"
	"time"
)

// Tile is the fundamental data object for this puzzle
// in the puzzle, it's a hex piece, but this is more general because that's just the way I thought to do it
type Tile struct {
	index int
	nbrs  map[int]bool
}

var (
	numIters     = flag.Int("num_iters", 19, "number of iterations of brownian motion to do")
	numTiles     = flag.Int("num_tiles", 74, "number of tiles to put in the grid")
	pcxn         = flag.Float64("pcxn", 0, "if non-zero, generates a random grid with pcxn probability of being connected to another neighbor")
	verbose      = flag.Bool("verbose", false, "show verbose log messages")
	printAsRat   = flag.Bool("print_as_rat", false, "prints the final probability as a floating point number")
	printTop     = flag.Int("print_top", 1, "number of tiles to print ordered from most probability, descending.")
)

var tile = make([]*Tile, 0)
var p = make([]*big.Rat, 0)

// Reset the state of the Tile Generator for hermetic Tests
func Reset() {
	tile = make([]*Tile, 0)
}

// AddNeighbors to the associated Tile
func (t *Tile) AddNeighbors(nbrs ...*Tile) {
	for _, n := range nbrs {
		t.nbrs[n.index] = true
		n.nbrs[t.index] = true
	}
}

// AddNeighborsByIndex instead of their object
func (t *Tile) AddNeighborsByIndex(idx ...int) {
	for _, i := range idx {
		t.nbrs[i] = true
		tile[i].nbrs[t.index] = true
	}
}

// NumNeighbors returns the number of neighbors for the associated Tile
func (t *Tile) NumNeighbors() int {
	return len(t.nbrs)
}

// GetNeighbors of the associated Tile
func (t *Tile) GetNeighbors() []*Tile {
	ret := []*Tile{}
	for idx := range t.nbrs {
		ret = append(ret, tile[idx])
	}
	return ret
}

// HasNeighbor returns true iff the the queried tile is a neighbor of the associated Tile
func (t *Tile) HasNeighbor(n *Tile) bool {
	_, ok := t.nbrs[n.index]
	return ok
}

func (t *Tile) String() string {
	return fmt.Sprintf("{index: %d, neighbors: %v}", t.index, t.nbrs)
}

// NewTile creates a new Tile and returns it.  index will increment every time it's called.
func NewTile() *Tile {
	ret := &Tile{
		index: len(tile),
		nbrs:  make(map[int]bool),
	}
	tile = append(tile, ret)
	return ret
}

// NewTiles creates a number of new tiles based on the parameter, n.
func NewTiles(n int) []*Tile {
	ret := []*Tile{}
	for i := 0; i < n; i++ {
		ret = append(ret, NewTile())
	}
	return ret
}

func main() {
	flag.Parse()
	NewTiles(*numTiles)

	// set up neighbors
	if *pcxn == 0 {
		tile[1].AddNeighborsByIndex(3, 4)
		tile[2].AddNeighborsByIndex(3, 5)
		tile[3].AddNeighborsByIndex(1, 2, 4, 5, 7)
		tile[4].AddNeighborsByIndex(1, 3, 6, 7, 9)
		tile[5].AddNeighborsByIndex(2, 3, 7, 10)
		tile[6].AddNeighborsByIndex(4, 9, 12)
		tile[7].AddNeighborsByIndex(3, 4, 5, 9, 10, 13)
		tile[8].AddNeighborsByIndex(11, 14)
		tile[9].AddNeighborsByIndex(4, 6, 7, 12, 13, 17)
		tile[10].AddNeighborsByIndex(5, 7, 13, 18)
		tile[11].AddNeighborsByIndex(8, 14, 15, 19)
		tile[12].AddNeighborsByIndex(6, 9, 16, 17, 22)
		tile[13].AddNeighborsByIndex(7, 9, 10, 17, 18, 23)
		tile[14].AddNeighborsByIndex(8, 11, 19)
		tile[15].AddNeighborsByIndex(11, 19, 20, 25)
		tile[16].AddNeighborsByIndex(12, 21, 22, 28)
		tile[17].AddNeighborsByIndex(9, 12, 13, 22, 23, 29)
		tile[18].AddNeighborsByIndex(10, 13, 23, 24, 30)
		tile[19].AddNeighborsByIndex(11, 14, 15, 25, 31)
		tile[20].AddNeighborsByIndex(15, 25, 26, 32)
		tile[21].AddNeighborsByIndex(16, 27, 28, 34)
		tile[22].AddNeighborsByIndex(12, 16, 17, 28, 29, 35)
		tile[23].AddNeighborsByIndex(13, 17, 18, 29, 30, 36)
		tile[24].AddNeighborsByIndex(18, 30)
		tile[25].AddNeighborsByIndex(15, 19, 20, 31, 32, 38)
		tile[26].AddNeighborsByIndex(20, 32, 33, 39)
		tile[27].AddNeighborsByIndex(21, 33, 34, 40)
		tile[28].AddNeighborsByIndex(16, 21, 22, 34, 35, 41)
		tile[29].AddNeighborsByIndex(17, 22, 23, 35, 36, 42)
		tile[30].AddNeighborsByIndex(18, 23, 24, 36)
		tile[31].AddNeighborsByIndex(19, 25, 37, 38, 43)
		tile[32].AddNeighborsByIndex(20, 25, 26, 38, 39, 44)
		tile[33].AddNeighborsByIndex(26, 27, 39, 40, 45)
		tile[34].AddNeighborsByIndex(21, 27, 28, 40, 41, 46)
		tile[35].AddNeighborsByIndex(22, 28, 29, 41, 42, 47)
		tile[36].AddNeighborsByIndex(23, 29, 30, 42)
		tile[37].AddNeighborsByIndex(31, 43, 48)
		tile[38].AddNeighborsByIndex(25, 31, 32, 43, 44, 49)
		tile[39].AddNeighborsByIndex(26, 32, 33, 44, 45, 50)
		tile[40].AddNeighborsByIndex(27, 33, 34, 45, 46, 51)
		tile[41].AddNeighborsByIndex(28, 34, 35, 46, 47, 52)
		tile[42].AddNeighborsByIndex(29, 35, 36, 47, 53)
		tile[43].AddNeighborsByIndex(31, 37, 38, 48, 49)
		tile[44].AddNeighborsByIndex(32, 38, 39, 49, 50, 55)
		tile[45].AddNeighborsByIndex(33, 39, 40, 50, 51, 56)
		tile[46].AddNeighborsByIndex(34, 40, 41, 51, 52, 57)
		tile[47].AddNeighborsByIndex(35, 41, 42, 52, 53, 58)
		tile[48].AddNeighborsByIndex(37, 43, 54)
		tile[49].AddNeighborsByIndex(38, 43, 44, 55)
		tile[50].AddNeighborsByIndex(39, 44, 45, 55, 56)
		tile[51].AddNeighborsByIndex(40, 45, 46, 56, 57, 59)
		tile[52].AddNeighborsByIndex(41, 46, 47, 57, 58, 60)
		tile[53].AddNeighborsByIndex(42, 47, 58, 61)
		tile[54].AddNeighborsByIndex(48)
		tile[55].AddNeighborsByIndex(44, 49, 50)
		tile[56].AddNeighborsByIndex(45, 50, 51, 59)
		tile[57].AddNeighborsByIndex(46, 51, 52, 59, 60, 62)
		tile[58].AddNeighborsByIndex(47, 52, 53, 60, 61, 63)
		tile[59].AddNeighborsByIndex(51, 56, 57, 62)
		tile[60].AddNeighborsByIndex(52, 57, 58, 62, 63, 64)
		tile[61].AddNeighborsByIndex(53, 58, 63)
		tile[62].AddNeighborsByIndex(57, 59, 60, 64, 65)
		tile[63].AddNeighborsByIndex(58, 60, 61, 64, 66)
		tile[64].AddNeighborsByIndex(60, 62, 63, 65, 66, 67)
		tile[65].AddNeighborsByIndex(62, 64, 67, 68)
		tile[66].AddNeighborsByIndex(63, 64, 67, 69)
		tile[67].AddNeighborsByIndex(64, 65, 66, 68, 69, 70)
		tile[68].AddNeighborsByIndex(65, 67, 70, 71)
		tile[69].AddNeighborsByIndex(66, 67, 70, 72)
		tile[70].AddNeighborsByIndex(67, 68, 69, 71, 72, 73)
		tile[71].AddNeighborsByIndex(68, 70, 73)
		tile[72].AddNeighborsByIndex(69, 70, 73)
		tile[73].AddNeighborsByIndex(70, 71, 72)
	} else {
		r := rand.New(rand.NewSource(time.Now().UnixNano()))
		for i := 0; i < len(tile); i++ {
			for j := i + 1; j < len(tile); j++ {
				if r.Float64() < *pcxn {
					tile[i].AddNeighborsByIndex(j)
				}
			}
		}
	}

	// set up probs
	pLast := make([]*big.Rat, len(tile))
	for i := 0; i < len(pLast); i++ {
		pLast[i] = big.NewRat(0, 1)
	}
	if *pcxn > 0 {
		pLast[0] = big.NewRat(1, 1)
	} else {
		pLast[1] = big.NewRat(1, 1)
	}
	pNext := make([]*big.Rat, len(tile))
	for i := 0; i < len(pNext); i++ {
		pNext[i] = big.NewRat(0, 1)
	}

	// loop through iterations and perform brownian motion, storing
	// probabilities at each iteration
	if *verbose {
		fmt.Println("iteration 0")
		fmt.Println(tile)
		fmt.Println(pLast)
		fmt.Println()
	}
	for t := 0; t < *numIters; t++ {
		for tIdx, t := range tile {
			pNext[tIdx] = big.NewRat(0, 1)
			for _, n := range t.GetNeighbors() {
				pOfN := big.NewRat(0, 1).Set(pLast[n.index])
				pNext[tIdx] = pNext[tIdx].Add(
					pNext[tIdx],
					pOfN.Mul(pOfN, big.NewRat(1, int64(n.NumNeighbors()))))
			}
		}
		// swap pLast and pNext
		for i := range pLast {
			t := pLast[i]
			pLast[i] = pNext[i]
			pNext[i] = t
		}
		if *verbose {
			fmt.Println("iteration ", t+1)
			fmt.Println(tile)
			fmt.Println(pLast)
			fmt.Println()
		}
	}

	// sanity check that all probs add up to 1
	s := big.NewRat(0, 1)
	for idx := range pLast {
		s.Add(s, pLast[idx])
	}
	if s.Cmp(big.NewRat(1, 1)) != 0 {
		fmt.Fprintf(os.Stderr, "All probabilities do not add up to 1.  Their sum is %v.", s)
		return
	}

	printed := make(map[int]bool)
	for i := 0; i < *printTop; i++ {
		// find most likely ending location
		maxP := big.NewRat(0, 1)
		maxI := -1
		for idx := range pLast {
			if maxP.Cmp(pLast[idx]) < 0 {
				if _, ok := printed[idx]; !ok {
					maxP = pLast[idx]
					maxI = idx
				}
			}
		}
		printed[maxI] = true

		if *printAsRat {
			fmt.Printf("Most Likely #%d: index %d with probability %v\n", i+1, maxI, maxP)
		} else {
			fmt.Printf("Most Likely #%d: index %d with probability %v\n", i+1, maxI, big.NewFloat(0).SetRat(maxP))
		}
	}
}