-module(huffman).
-export([print_codewords/1]).

print_codewords(Text) ->
    Tree = mk_tree(freq_dict(Text)),
    Codewords = codewords(Tree),
    lists:foreach(fun({Chr, Code}) -> io:fwrite("~c: ~w~n", [Chr, Code]) end, Codewords).

codewords({L, R}) ->
    codewords(L, <<0:1>>) ++ codewords(R, <<1:1>>).

codewords({L, R}, <<Bits/bits>>) ->
    codewords(L, <<Bits/bits, 0:1>>) ++ codewords(R, <<Bits/bits, 1:1>>);
codewords(Symbol, <<Bits/bitstring>>) ->
    [{Symbol, Bits}].

mk_tree([{Tmp, _}]) -> Tmp;
mk_tree(Tmp) ->
    [{A1, B1}, {A2, B2} | Rest] = lists:keysort(2, Tmp),
    mk_tree([{{A1, A2}, B1+B2} | Rest]).

freq_dict(Text) ->
    freq_dict(lists:sort(Text), []).

freq_dict([], Acc) -> Acc;
freq_dict([Chr | Rest], Acc) ->
    {ChrBlock, OtherChrs} = lists:splitwith(fun(X) -> X == Chr end, Rest),
    freq_dict(OtherChrs, [{Chr, 1+length(ChrBlock)} | Acc]).