package main

import (
	"container/heap"
	"container/list"
	"fmt"
	"log"
)

type charFreq struct {
	letter uint8
	count  int
}

type huffmanNode struct {
	letter uint8
	freq   float32
	left   *huffmanNode
	right  *huffmanNode
}

func encodeString(inputString string, rootNode *huffmanNode) string {
	if len(inputString) == 0 {
		log.Fatal("Empty inputString")
	}
	if rootNode == nil {
		log.Fatal("Nil rootNode")
	}

	encodingCache := make(map[uint8]string)
	var outputString string
	for _, c := range inputString {
		encodedChar, isPresent := encodingCache[uint8(c)]
		if !isPresent {
			encodedChar = getEncodingForChar(uint8(c), rootNode, "")
			encodingCache[uint8(c)] = encodedChar
		}
        outputString = outputString + encodedChar
    }

	return outputString
}

func getEncodingForChar(letter uint8, huffmanTree *huffmanNode, output string) string {
	if letter == huffmanTree.letter {
		return output
	} else {
		var result string
		if huffmanTree.left != nil {
			result = getEncodingForChar(letter, huffmanTree.left, output+"0")
		}
		if huffmanTree.right != nil && result == "" {
			result = getEncodingForChar(letter, huffmanTree.right, output+"1")
		}
		return result
	}
}

func buildHuffmanTree(inputStrLength int, charFreqList *list.List) *huffmanNode {
	pq := make(PriorityQueue, charFreqList.Len())

	for i, e := 0, charFreqList.Front(); e != nil; i, e = i+1, e.Next() {
		charFreq1 := e.Value.(charFreq)
		pq[i] = &Item{
			priority: float32(charFreq1.count) / float32(inputStrLength),
			value: &huffmanNode{
				letter: charFreq1.letter,
				freq:   float32(charFreq1.count) / float32(inputStrLength),
			},
		}
	}
	heap.Init(&pq) //insert entire array into heap

	for pq.Len() > 1 {
		x := heap.Pop(&pq).(*Item)
		y := heap.Pop(&pq).(*Item)
		z := Item{
			priority: x.priority + y.priority,
			value: &huffmanNode{
				freq:  x.priority + y.priority,
				left:  x.value,
				right: y.value,
			},
		}
		heap.Push(&pq, &z)
	}

	return heap.Pop(&pq).(*Item).value
}

func determineCharFreq(inputStr string) *list.List {
	charFreqList := list.New()
	for i := 0; i < len(inputStr); i++ {
		for e := charFreqList.Front(); true; e = e.Next() {
			if e == nil {
				charFreqList.PushBack(charFreq{letter: inputStr[i], count: 1})
				break
			}
			charFreq1 := e.Value.(charFreq) //this makes a copy(!)
			if charFreq1.letter == inputStr[i] {
				e.Value = charFreq{letter: inputStr[i], count: charFreq1.count + 1}
				break
			}
		}
	}

	for e := charFreqList.Front(); e != nil; e = e.Next() {
		charFreq1 := e.Value.(charFreq)
		fmt.Printf("%c=%v\n", charFreq1.letter, charFreq1.count)
	}

	return charFreqList
}

func main() {
	inputStr := "TAATTAGAAATTCTATTATA"
	fmt.Println(inputStr)
	charFreqList := determineCharFreq(inputStr)
	rootNode := buildHuffmanTree(len(inputStr), charFreqList)
	fmt.Println("output:", encodeString(inputStr, rootNode))
}