module Huffman where

import qualified Data.Map.Strict as Map
import Auxiliaries ( Bit (..), FileContent (FileContent, fc_codingTable, fc_content), Bitlist (Bitlist), binaryToFile, binaryFromFile)
import Data.Maybe (fromMaybe)
import Data.Char (toUpper)
import qualified Data.IntMap.Strict as IntMap
import System.IO ( hGetContents, hPutStr, withFile, IOMode(WriteMode, ReadMode) )

type IntMap = IntMap.IntMap
type Map = Map.Map
type CodingTable = Map Char [Bit]

-----------------------------------------------------------------------
-- Aufgabe 2
-----------------------------------------------------------------------

data Node = Leaf Char Int | Inner Node Node Int deriving (Show)
newtype HTree = Root Node deriving (Show)

getFrequency :: HTree -> Int
getFrequency (Root tree) =
    case tree of
        Leaf _ frequency -> frequency
        Inner _ _ frequency -> frequency

getNode :: HTree -> Node
getNode (Root node) = node

exercise1TreeValid :: HTree
exercise1TreeValid = Root (Inner (Inner (Leaf 'E' 158) (Inner (Leaf 'N' 97) (Leaf 'I' 82) 179) 337) (Inner (Inner (Leaf 'R' 77) (Leaf 'S' 67) 144) (Inner (Leaf 'T' 64) (Leaf 'A' 61) 125) 269) 606)

exercise1TreeInvalid :: HTree
exercise1TreeInvalid = Root (Inner (Leaf 'A' 21) (Leaf 'B' 20) 42)

-----------------------------------------------------------------------
-- Aufgabe 3
-----------------------------------------------------------------------

isConsistent :: HTree -> Bool
isConsistent (Root node) =
    case node of
        Leaf _ _ -> True
        Inner left right frequency ->
            case left of
                Leaf _ leftFrequency ->
                    case right of
                        Leaf _ rightFrequency -> frequency == leftFrequency + rightFrequency
                        Inner _ _ rightFrequency -> frequency == leftFrequency + rightFrequency && isConsistent (Root right)
                Inner _ _ leftFrequency ->
                    case right of
                        Leaf _ rightFrequency -> frequency == leftFrequency + rightFrequency && isConsistent (Root left)
                        Inner _ _ rightFrequency -> frequency == leftFrequency + rightFrequency && isConsistent (Root left) && isConsistent (Root right)

-----------------------------------------------------------------------
-- Aufgabe 4
-----------------------------------------------------------------------

data Direction = Left | Right
type BitList = [Bit]

toCodingTable :: HTree -> CodingTable
toCodingTable (Root node) =
    case node of
        Leaf char _ -> Map.singleton char [Zero]
        Inner left right _ -> Map.union (Map.fromList (toCodingTableList (Root left) Huffman.Left)) (Map.fromList (toCodingTableList (Root right) Huffman.Right)) 

toCodingTableList :: HTree -> Direction -> [(Char,BitList)]
toCodingTableList (Root node) direction =
    case direction of
        Huffman.Left ->
            case node of
                Inner left right _ -> prependToBitLists Zero (toCodingTableList (Root left) Huffman.Left) ++ prependToBitLists Zero (toCodingTableList (Root right) Huffman.Right)
                Leaf char _ -> prependToBitLists Zero [(char,[])]
        Huffman.Right ->
            case node of
                Inner left right _ -> prependToBitLists One (toCodingTableList (Root left) Huffman.Left) ++ prependToBitLists One (toCodingTableList (Root right) Huffman.Right)
                Leaf char _ -> prependToBitLists One [(char,[])]

-- Prepends given Bit in front of each BitList inside the tuples of given list.
prependToBitLists :: Bit -> [(Char, BitList)] -> [(Char, BitList)]
prependToBitLists bitToPrepend = Prelude.map (\(char, bitList) -> (char, bitToPrepend : bitList))

-----------------------------------------------------------------------
-- Aufgabe 4 Test Data
-----------------------------------------------------------------------

treeValidCodingTable :: CodingTable
treeValidCodingTable = Map.fromList [('E',[Zero,Zero]), ('N',[Zero,One,Zero]), ('I',[Zero,One,One]), ('R',[One,Zero,Zero]), ('S',[One,Zero,One]), ('T',[One,One,Zero]), ('A',[One,One,One])]

treeOnlyLeafCodingTable :: CodingTable
treeOnlyLeafCodingTable = Map.fromList [('A',[Zero])]

-----------------------------------------------------------------------
-- Aufgabe 5
-----------------------------------------------------------------------

encode :: CodingTable -> String -> [Bit]
encode table input =
    case input of
        (firstChar:rest) -> encodeChar table firstChar ++ encode table rest
        [] -> []

encodeChar :: CodingTable -> Char -> [Bit]
encodeChar table char = fromMaybe [] (Map.lookup char table)

-----------------------------------------------------------------------
-- Aufgabe 5 Test Data
-----------------------------------------------------------------------

codingTableForTesting :: CodingTable
codingTableForTesting = toCodingTable exercise1TreeValid

enteEncoded :: BitList
enteEncoded = [Zero,Zero, Zero,One,Zero, One,One,Zero, Zero,Zero]

erenEncoded :: BitList
erenEncoded = [Zero,Zero, One,Zero,Zero, Zero,Zero, Zero,One,Zero]

startEncoded :: BitList
startEncoded = [One,Zero,One, One,One,Zero, One,One,One, One,Zero,Zero, One,One,Zero]

rainEncoded :: BitList
rainEncoded = [One,Zero,Zero, One,One,One, Zero,One,One, Zero,One,Zero]

-----------------------------------------------------------------------
-- Aufgabe 6
-----------------------------------------------------------------------

decode :: HTree -> [Bit] -> String
decode tree encodedString =
    case encodedString of
        [] -> []
        _ -> 
            let (steps, char) = walkTillLeaf tree encodedString 0
                in char ++ decode tree (drop steps encodedString)

walkTillLeaf :: HTree -> [Bit] -> Int -> (Int, String)
walkTillLeaf (Root tree) directions step =
    case tree of
        Leaf char _ -> (step, [char])
        Inner left right _ ->
            case directions of
                (first:rest) ->
                    case first of
                        Zero -> walkTillLeaf (Root left) rest (step + 1)
                        One -> walkTillLeaf (Root right) rest (step + 1)
                [] -> (0,"")

-----------------------------------------------------------------------
-- Aufgabe 7
-----------------------------------------------------------------------

instance Eq Node where
    (Leaf char1 frequency1) == (Leaf char2 frequency2) = toUpper char1 == toUpper char2 && frequency1 == frequency2
    (Leaf _ _) == (Inner {}) = False
    (Inner {}) == (Leaf _ _) = False
    (Inner node11 node12 frequency1) == (Inner node21 node22 frequency2) = node11 == node21 && node12 == node22 && frequency1 == frequency2

instance Ord Node where
    (Leaf _ frequency1) >= (Leaf _ frequency2) = frequency1 >= frequency2
    (Leaf _ frequency1) >= (Inner _ _ frequency2) = frequency1 >= frequency2
    (Inner _ _ frequency1) >= (Leaf _ frequency2) = frequency1 >= frequency2
    (Inner _ _ frequency1) >= (Inner _ _ frequency2) = frequency1 >= frequency2
    (Leaf _ frequency1) <= (Leaf _ frequency2) = frequency1 <= frequency2
    (Leaf _ frequency1) <= (Inner _ _ frequency2) = frequency1 <= frequency2
    (Inner _ _ frequency1) <= (Leaf _ frequency2) = frequency1 <= frequency2
    (Inner _ _ frequency1) <= (Inner _ _ frequency2) = frequency1 <= frequency2

instance Eq HTree where
    (Root node1) == (Root node2) = node1 == node2

buildHTree :: String -> HTree
buildHTree input =
    let charFrequencyMap = combineFrequencies (addOccurences input)
        in toHTree (toHTreeIntMap (toHTreeNodes charFrequencyMap))

--map ['r','a','i','n','a'] becomes [('r',1),('a',1),('i',1),('n',1),('a',1)]
addOccurences :: String -> [(Char,Int)]
addOccurences = Prelude.map (\char -> (char,1))

--fromListWith (+) [('r',1),('a',1),('i',1),('n',1),('a',1)] becomes [('r',1),('a',2),('i',1),('n',1)]
combineFrequencies :: [(Char,Int)] -> Map Char Int
combineFrequencies = Map.fromListWith (+)

toHTreeNodes :: Map Char Int -> [(Int, [HTree])]
toHTreeNodes frequenciesMap = 
    let charFrequenciesList = Map.toList frequenciesMap
        in Prelude.map (\(char,frequency) -> (frequency, [Root (Leaf char frequency)])) charFrequenciesList

toHTreeIntMap :: [(Int, [HTree])] -> IntMap [HTree]
toHTreeIntMap = IntMap.fromListWith (++)

intMapAaron :: IntMap [HTree]
intMapAaron = IntMap.fromListWith (++) [(1,[Root (Leaf 'A' 1)]),(1,[Root (Leaf 'a' 1)]),(1,[Root (Leaf 'r' 1)]),(1,[Root (Leaf 'o' 1)]),(1,[Root (Leaf 'n' 1)])]

intMapB :: IntMap [HTree]
intMapB = IntMap.fromListWith (++) [(3,[Root (Leaf 'A' 3)]),(1,[Root (Leaf 'a' 1)]),(1,[Root (Leaf 'r' 1)]),(2,[Root (Leaf 'o' 2)]),(4,[Root (Leaf 'n' 4)])]

toHTree :: IntMap [HTree] -> HTree
toHTree map =
    case IntMap.minView map of
        Just (minimums, newMap) ->
            case minimums of
                [] -> 
                    if IntMap.size newMap > 0
                        then toHTree newMap
                        else (Root (Leaf 'A' 0))
                [x] ->
                    case IntMap.minView newMap of
                        Just (minimums2, newMap2) ->
                            case minimums2 of
                                [] -> x
                                (y:rest) -> toHTree (IntMap.insertWith (++) (getFrequency x + getFrequency y) [Root (Inner (getNode x) (getNode y) (getFrequency x + getFrequency y))] (IntMap.insertWith (++) (getFrequency y) rest newMap2))
                        _ -> x
                (x:y:rest) ->
                    toHTree (IntMap.insertWith (++) (getFrequency x + getFrequency y) [Root (Inner (getNode x) (getNode y) (getFrequency x + getFrequency y))] (IntMap.insertWith (++) (getFrequency x) rest newMap))
        _ -> case IntMap.maxView map of
            Just (maximums, _) -> head maximums
            _ -> (Root (Leaf ' ' 0))

-----------------------------------------------------------------------
-- Aufgabe 7 Test
-----------------------------------------------------------------------

treeAaron :: HTree
treeAaron = Root (Inner (Inner (Leaf 'r' 1) (Leaf 'o' 1) 2) (Inner (Leaf 'A' 1) (Inner (Leaf 'n' 1) (Leaf 'a' 1) 2) 3) 5)

treeMISSISSIPPI :: HTree
treeMISSISSIPPI = Root (Inner (Leaf 'I' 4) (Inner (Inner (Leaf 'M' 1) (Leaf 'P' 2) 3) (Leaf 'S' 4) 7) 11)

treeAAABC :: HTree
treeAAABC = Root (Inner (Inner (Leaf 'C' 1) (Leaf 'B' 1) 2) (Leaf 'A' 3) 5)

-----------------------------------------------------------------------
-- Aufgabe 8
-----------------------------------------------------------------------

toDecodeTree :: CodingTable -> HTree
toDecodeTree table = foldl1 mergeTrees (map (\(char, bitCode) -> partTree char bitCode) (Map.toList table))

partTree :: Char -> [Bit] -> HTree
partTree char bits = 
    case bits of
        [] -> Root (Leaf char 0)
        (bit:rest) ->
            let next = partTree char rest
                in case bit of
                    Zero -> Root (Inner (getNode next) (Leaf 'e' 1) 0)
                    One -> Root (Inner (Leaf 'e' 1) (getNode next) 0)

mergeTrees :: HTree -> HTree -> HTree
mergeTrees (Root (Leaf char1 0)) (Root (Leaf char2 0)) = Root (Inner (Leaf char1 0) (Leaf char2 0) 0)
mergeTrees (Root (Leaf char 0)) (Root (Leaf 'e' 1)) = Root (Leaf char 0)
mergeTrees (Root (Leaf 'e' 1)) (Root (Leaf char 0)) = Root (Leaf char 0)
mergeTrees (Root (Inner left1 right1 _)) (Root (Inner left2 right2 _)) = Root (Inner (getNode (mergeTrees (Root left1) (Root left2))) (getNode (mergeTrees (Root right1) (Root right2))) 0)
mergeTrees (Root (Leaf _ _)) right = right
mergeTrees left (Root (Leaf _ _)) = left

-----------------------------------------------------------------------
-- Aufgabe 8 Test
-----------------------------------------------------------------------

testCodingTable :: CodingTable
testCodingTable = Map.fromList [('E', [Zero,Zero]), ('N', [Zero,One,Zero]), ('I', [Zero,One,One]), ('R', [One,Zero,Zero]), ('S', [One,Zero,One]), ('T', [One,One,Zero]), ('A', [One,One,One])]

testCodingTableHTree :: HTree
testCodingTableHTree = Root (Inner (Inner (Leaf 'E' 0) (Inner (Leaf 'N' 0) (Leaf 'I' 0) 0) 0) (Inner (Inner (Leaf 'R' 0) (Leaf 'S' 0) 0) (Inner (Leaf 'T' 0) (Leaf 'A' 0) 0) 0) 0)

-----------------------------------------------------------------------
-- Aufgabe 9
-----------------------------------------------------------------------

encodeFile :: FilePath -> FilePath -> IO ()
encodeFile input output =
    withFile input ReadMode $ \handle -> do
    inputContent <- hGetContents handle
    let htree = buildHTree inputContent
        codingTable = toCodingTable htree
        encodedContent = encode codingTable inputContent
        fileContent = FileContent {fc_codingTable = (map (\(char, bitList) -> (char, Bitlist bitList)) (Map.toList codingTable)), fc_content = (Bitlist encodedContent)}
        in do binaryToFile output fileContent

decodeFile :: FilePath -> FilePath -> IO ()
decodeFile input output = do
    fileContent <- binaryFromFile input
    let codingTable = fc_codingTable fileContent
        encodedContent = fc_content fileContent
        decodeTree = toDecodeTree (Map.fromList (map (\(char, Bitlist bitList) -> (char, bitList)) codingTable))
        decodedString = decode decodeTree (bitlistToList encodedContent)
        in do withFile output WriteMode (\handle -> do hPutStr handle decodedString)

bitlistToList :: Bitlist -> [Bit]
bitlistToList (Bitlist bits) = bits

-----------------------------------------------------------------------
-- Aufgabe 9
-----------------------------------------------------------------------
-- siehe Tests.hs

-----------------------------------------------------------------------
-- Aufgabe 10
-----------------------------------------------------------------------
-- siehe Main.hs