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
import Data.List (maximumBy, minimumBy)
import Text.Printf (printf)

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)

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)) (Root (Leaf 'A' 0))

--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 (++)

toHTree :: IntMap [HTree] -> HTree -> HTree
toHTree map lastTree =
    let Root lastTreeRoot = lastTree
        in case lastTreeRoot of
            Leaf _ lastTreeRootFrequency -> 
                case IntMap.minView map of
                    Just (minimums, newMap) ->
                        case minimums of
                            [] -> toHTree newMap lastTree
                            [x] ->
                                let Root xRoot = x
                                    in case xRoot of
                                        Leaf _ xFrequency ->
                                            if lastTreeRootFrequency == 0
                                                then toHTree newMap x
                                                else if lastTreeRootFrequency < xFrequency
                                                    then toHTree newMap (Root (Inner lastTreeRoot xRoot (lastTreeRootFrequency + xFrequency)))
                                                    else toHTree newMap (Root (Inner xRoot lastTreeRoot (lastTreeRootFrequency + xFrequency)))
                                        Inner _ _ xFrequency ->
                                            if lastTreeRootFrequency == 0
                                                then toHTree newMap x
                                                else if lastTreeRootFrequency < xFrequency
                                                    then toHTree newMap (Root (Inner lastTreeRoot xRoot (lastTreeRootFrequency + xFrequency)))
                                                    else toHTree newMap (Root (Inner xRoot lastTreeRoot (lastTreeRootFrequency + xFrequency)))
                            (x:y:rest) ->
                                let Root xRoot = x
                                    Root yRoot = y
                                    in case xRoot of
                                        Leaf _ xFrequency ->
                                            if lastTreeRootFrequency == 0
                                                then toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner xRoot yRoot (xFrequency * 2)))
                                                else if lastTreeRootFrequency < xFrequency
                                                    then toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner lastTreeRoot (Inner xRoot yRoot (xFrequency * 2)) (lastTreeRootFrequency + (xFrequency * 2))))
                                                    else toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner (Inner xRoot yRoot (xFrequency * 2)) lastTreeRoot (lastTreeRootFrequency + (xFrequency * 2))))
                                        Inner _ _ xFrequency ->
                                            if lastTreeRootFrequency == 0
                                                then toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner xRoot yRoot (xFrequency * 2)))
                                                else if lastTreeRootFrequency < xFrequency
                                                    then toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner lastTreeRoot (Inner xRoot yRoot (xFrequency * 2)) (lastTreeRootFrequency + (xFrequency * 2))))
                                                    else toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner (Inner xRoot yRoot (xFrequency * 2)) lastTreeRoot (lastTreeRootFrequency + (xFrequency * 2))))
                    _ -> lastTree
            Inner _ _ lastTreeRootFrequency ->
                case IntMap.minView map of
                    Just (minimums, newMap) ->
                        case minimums of
                            [] -> toHTree newMap lastTree
                            [x] ->
                                let Root xRoot = x
                                    in case xRoot of
                                        Leaf _ xFrequency ->
                                            if lastTreeRootFrequency < xFrequency
                                                then toHTree newMap (Root (Inner lastTreeRoot xRoot (lastTreeRootFrequency + xFrequency)))
                                                else toHTree newMap (Root (Inner xRoot lastTreeRoot (lastTreeRootFrequency + xFrequency)))
                                        Inner _ _ xFrequency ->
                                            if lastTreeRootFrequency < xFrequency
                                                then toHTree newMap (Root (Inner lastTreeRoot xRoot (lastTreeRootFrequency + xFrequency)))
                                                else toHTree newMap (Root (Inner xRoot lastTreeRoot (lastTreeRootFrequency + xFrequency)))
                            (x:y:rest) ->
                                let Root xRoot = x
                                    Root yRoot = y
                                    in case xRoot of
                                        Leaf _ xFrequency ->
                                            if lastTreeRootFrequency < xFrequency
                                                then toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner lastTreeRoot (Inner xRoot yRoot (xFrequency * 2)) (lastTreeRootFrequency + (xFrequency * 2))))
                                                else toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner (Inner xRoot yRoot (xFrequency * 2)) lastTreeRoot (lastTreeRootFrequency + (xFrequency * 2))))
                                        Inner _ _ xFrequency ->
                                            if lastTreeRootFrequency < xFrequency
                                                then toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner lastTreeRoot (Inner xRoot yRoot (xFrequency * 2)) (lastTreeRootFrequency + (xFrequency * 2))))
                                                else toHTree (IntMap.insert xFrequency rest newMap) (Root (Inner (Inner xRoot yRoot (xFrequency * 2)) lastTreeRoot (lastTreeRootFrequency + (xFrequency * 2))))
                    _ -> lastTree

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

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

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

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

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

--type CodingTable = Map Char [Bit]
type Depth = Int

toDecodeTree :: CodingTable -> HTree
toDecodeTree table =
    let listWithDepth = (codingListToCodingListWithDepth (codingTableToList table))
        listWithDepthInverted = invertDepths listWithDepth (getMinimumDepth listWithDepth) (getMaximumDepth listWithDepth)
        in buildHTree (depthListToString listWithDepthInverted)

depthListToString :: [(Char, Depth)] -> String
depthListToString list =
    case list of
        [] -> ""
        ((char, depth):rest) ->  (replicate depth char) ++ depthListToString rest

invertDepths :: [(Char, Depth)] -> Depth -> Depth -> [(Char, Depth)]
invertDepths list minDepth maxDepth = map (\(char, depth) -> (char, (minDepth + maxDepth - depth))) list

codingTableToList :: CodingTable -> [(Char, [Bit])]
codingTableToList table = Map.toList table

codingListToCodingListWithDepth :: [(Char, [Bit])] -> [(Char, Depth)]
codingListToCodingListWithDepth tableList =
    map (\(char, bitList) -> (char, length bitList)) tableList

getMaximumDepth :: [(Char, Depth)] -> Depth
getMaximumDepth = (\(_, d) -> d) . maximumBy compareByDepth

getMinimumDepth :: [(Char, Depth)] -> Depth
getMinimumDepth = (\(_, d) -> d) . minimumBy compareByDepth

compareByDepth :: (Char, Depth) -> (Char, Depth) -> Ordering
compareByDepth (_, depth1) (_, depth2)
    | depth1 < depth2 = LT
    | depth1 > depth2 = GT
    | otherwise = EQ

-----------------------------------------------------------------------
-- 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 (Leaf 'E' 3) (Inner (Inner (Leaf 'I' 2) (Leaf 'A' 2) 4) (Inner (Inner (Leaf 'R' 2) (Leaf 'N' 2) 4) (Inner (Leaf 'T' 2) (Leaf 'S' 2) 4) 8) 12) 15)

-----------------------------------------------------------------------
-- 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
            printf ("\n" ++ show htree ++ "\n")
            --printf (show codingTable ++ "\n")
            --printf (show encodedContent ++ "\n")
            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
            printf ("\n" ++ show decodeTree)
            withFile output WriteMode (\handle -> do hPutStr handle decodedString)

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