package main

import (
  "fmt"
)


type IndexedSet struct {
  v,i int
}

var it = 0

func fillByFormule( slc *[]IndexedSet, form func (int) int) {
  it++
  for n, i := 0, 0 ; n < 10000 ; i++ {
    if n > 999 && n < 10000 {
      if n % 100 > 9 {
        *slc = append(*slc, IndexedSet{v:n, i:it})
      }
    }
    n = form(i)
  }
}
func fillByNumber( dict map[int][]IndexedSet, num []IndexedSet) {
  for _,n := range num {
    dict[n.v/100] = append(dict[n.v/100], n)
  }
}

func containsIndex( ref []IndexedSet, index int) bool {
  for _,n := range ref {
    if n.i == index {
      return true
    }
  }
  return false
}
func containsValue( ref []IndexedSet, value int) bool {
  for _,n := range ref {
    if n.v == value {
      return true
    }
  }
  return false
}

func main () {
  trin := []IndexedSet{}
  sqrt := []IndexedSet{}
  pent := []IndexedSet{}
  hexa := []IndexedSet{}
  hept := []IndexedSet{}
  octa := []IndexedSet{}

  numMap := make(map[int][]IndexedSet)


  fillByFormule(&trin, func(x int) int { return (x*(x+1))>>1 })
  fillByFormule(&sqrt, func(x int) int { return x*x })
  fillByFormule(&pent, func(x int) int { return (x*(3*x-1))>>1 })
  fillByFormule(&hexa, func(x int) int { return x*(2*x-1) })
  fillByFormule(&hept, func(x int) int { return (x*(5*x-3))>>1 })
  fillByFormule(&octa, func(x int) int { return x*(3*x-2) })

  fillByNumber(numMap, trin)
  fillByNumber(numMap, sqrt)
  fillByNumber(numMap, pent)
  fillByNumber(numMap, hexa)
  fillByNumber(numMap, hept)

  pile := []IndexedSet{}

  pile = append( pile, octa...)


  cycle := []IndexedSet{}

  for len(pile) > 0 {
    if len(cycle) >= 6 {
      if (cycle[0].v/100) == (cycle[5].v%100) {
        fmt.Println(cycle)
        sum := 0
        for _,v := range cycle {
          sum += v.v
        }
        fmt.Println(sum)
      }
      for pile[len(pile)-1].i != -1 {
        pile = pile[:len(pile)-1]
      }
    }


    n := pile[len(pile)-1]
    pile = pile[:len(pile)-1]

    if n.i == -1 {
      if len(cycle) > 0 {
        cycle = cycle[:len(cycle)-1]
      } else {
        fmt.Println("Houston, we got a problem")
      }
      continue
    }

    n_2 := n.v % 100

    if !containsIndex(cycle, n.i) && !containsValue( cycle, n.v) {
      cycle = append(cycle, n)
      pile = append(pile, IndexedSet{ v:-1, i: -1} )
      if vList,ok := numMap[n_2] ; ok && vList != nil {
        for _,v := range vList {
          if !containsIndex( cycle, v.i ) {
            pile = append(pile, v)
          }
        }
      }
    }
  }
}