from heapq import heappush, heappop, heapify
from collections import defaultdict

def create_prob_model(message):
    """
    Create an alphabeth based on the symbols in the message, and
    attach a frequency to each symbol
    """
    
    prob_model = defaultdict(int) # Every new key is given a default value of 0
    for char in message:
        prob_model[char] += 1 / len(message)
        
    return prob_model

def huff_encode(message):
    """Huffman encoding of the input message.
    
    Args:
        message: String, a sequence of characters (symbols)
                          
    Returns:
        source_code: Dict where every entry is {symbol: (frequency, codeword)}
        encoded_message: String, bitstream of the encoded input message using the source_code.
    
    """
    
    prob_model = create_prob_model(message)
    
    heap = [[freq, [sym, ""]] for sym, freq in prob_model.items()]
    heapify(heap)
    while len(heap) > 1:
        lo = heappop(heap)
        hi = heappop(heap)
        for pair in lo[1:]:
            pair[1] = '0' + pair[1]
        for pair in hi[1:]:
            pair[1] = '1' + pair[1]
        heappush(heap, [lo[0] + hi[0]] + lo[1:] + hi[1:])
    symbol_code_pairs = sorted(heappop(heap)[1:], key=lambda p: (len(p[-1]), p))
    
    source_code = {}
    for pair in symbol_code_pairs:
        symbol = pair[0]
        codeword = pair[1]
        freq = prob_model[symbol]
        source_code[symbol] = (freq, codeword)
    
    encoded_message = ''
    for symbol in message:
        encoded_message += source_code[symbol][1]
    
    return source_code, encoded_message

message = "this is an example for huffman encoding"
source_code, encoded_message = huff_encode(message)

print("Message:         {}".format(message))
print("Encoded message: {}\n".format(encoded_message))

print("Symbol Frequency Codeword")
for symbol, pair in sorted(source_code.items()):
    print("{0:>6} {1:>9,.2f} {2:>8}".format(symbol, pair[0], pair[1]))

import numpy as np

def scale_interval_edges(orig_interval_edges, min_val, max_val):
    """Scale the original interval edges to the new min and max values"""
    new_edges = min_val + (max_val - min_val)*orig_interval_edges
    return new_edges

def symbol_to_index(symbol, alphabeth):
    """Map symbol to index according to alphabeth (1-based indexing since we are working on intervals,
       the returned index will be according to the upper edge of the interval.)"""
    assert len(set(alphabeth)) == len(alphabeth), 'Redundant alphabeth'
    assert symbol in alphabeth, 'Symbol {} not in alphabeth'.format(symbol)
    return alphabeth.index(symbol) + 1 #

def get_interval_from_symbol(current_symbol, alphabeth, current_min, current_max, orig_interval_edges):
    """Get a new interval for the new symbol, based on the current min and max and original interval edges"""
    curr_signal_ind = symbol_to_index(current_symbol, alphabeth)
    curr_interval_edges = scale_interval_edges(orig_interval_edges, current_min, current_max)
    new_min = curr_interval_edges[curr_signal_ind - 1]
    new_max = curr_interval_edges[curr_signal_ind]
    
    return (new_min, new_max)

def interval_edges(pmf):
    """Return a list of interval edges, given the probability mass function"""
    return np.array([np.sum(pmf[:i]) for i in range(len(pmf) + 1)])

def arithmetic_intervals(alphabeth, signal, pmf):
    """Get a list of intervals obtained by arithmetic coding of signal.
    
    Args:
        alphabeth: List of unique characters.
        signal: String of symbols, where every symbol need to be a member of alphabeth.
        pmf: Probability mass function for elements in alphabeth. The order follows
             the order of alphabeth, that is pmf[i] is the probability of alphabeth[i].
           
    Returns:
        intervals: List of (min, max) float tuples, where (min, max) is the interval
                   edges of the respective symbol. This list has one interval per symbol in
                   the input signal.
    """
    orig_interval_edges = interval_edges(pmf)
    signal_list = list(signal)
    curr_min, curr_max = get_interval_from_symbol(signal_list[0], alphabeth, 0.0, 1.0, orig_interval_edges)
    intervals = [(curr_min, curr_max)]
    
    for i, symbol in enumerate(signal_list[1:]):
        curr_min, curr_max = get_interval_from_symbol(symbol, alphabeth, curr_min, curr_max, orig_interval_edges)
        intervals.append((curr_min, curr_max))
    return intervals

def binary_sequence(decimal):
    """Return a binary sequence representation of input.
    
    Args:
        decimal: scalar float in [0, 1) # Important that the interval is half-open
    
    Returns:
        bin_sequence: List of integers in {0, 1}
    """
    reminder = decimal
    bin_sequence = [int(np.floor(2*reminder))]
    ind = 0
    while reminder > 1e-10: # This is not very robust w.r.t. float precision
        reminder = 2*reminder - bin_sequence[ind]
        bin_sequence.append(int(np.floor(2*reminder)))
        ind += 1
    return bin_sequence

def shortest_binary(d_interval):
    """Find shortest binary representation within a decimal interval.
    
    Args:
        d_interval: Tuple (d_min, d_max) of floats. This is a
                    half open interval [d_min, d_max) ⊆ [0, 1).
        
    Returns:
        bin_seq: Binary representation of shortest length within d_interval.
    """
    d_min, d_max = d_interval
    assert d_min < d_max, 'Need strictly increasing interval'
    assert d_min >= 0, 'Negative lower bound on interval'
    assert d_max < 1, 'Upper interval bound greater or equal to 1'
    
    # Initialize current interval (c_min, c_max)
    c_min = 0.0
    c_max = 1.0
    
    # We continue until the current interval (c_min, c_max) is contained inside
    # the given interval (d_min, d_max).
    
    # Note: Adding a 1 to the binary sequence increases c_min, but leaves c_max fixed.
    #       Adding a 0 to the binary sequence decreases c_max, but leaves c_min fixed.
    
    k = 1
    bin_seq = []
    while True:
        # Allways add 1 to bin_seq if possible
        if c_min < d_min and c_min + 1/2**k < d_max:
            c_min = c_min + 1/2**k
            bin_seq.append(1)
        else:
            if c_max > d_max and c_max - 1/2**k > d_min:
                c_max = c_max - 1/2**k
                bin_seq.append(0)
            else:
                # No change is made to the current interval, we are finished
                break
        k = k + 1
    return bin_seq            

def arithmetic_encoding(alphabeth, pmf, signal):
    """Encode a signal to a binary sequence.
    
    Args:
        alphabeth: List of unique characters.
        signal: String of symbols, where every symbol need to be a member of alphabeth.
        pmf: Probability mass function for elements in alphabeth. The order follows
             the order of alphabeth, that is pmf[i] is the probability of alphabeth[i].
    
    Returns: 
        bin_seq: Arithmetic encoding of input signal.
    """
    intervals = arithmetic_intervals(alphabeth, signal, pmf)
    bin_seq = shortest_binary(intervals[-1])
    return bin_seq

def interval_in_increasing(number, increasing_list):
    """Find the interval where number is located within the list of increasing numbers."""
    assert number >= increasing_list[0], 'Number is outside increasing_list'
    assert number <= increasing_list[-1], 'Number is outside increasing_list'
    interval = None
    for i in range(1, len(increasing_list)):
        if number >= increasing_list[i-1] and number <= increasing_list[i]:
            interval = (i-1, i)
    assert interval is not None, 'Invalid input to interval_in_increasing'
    return interval

def decimal_from_binary(bin_seq):
    """Every decimal number d_10 in [0, 1) can be represented as a power series
    
        d_10 = b_1*1/2^1 + b_2*1/2^2 + b_3*1/2^3 + ...
        
        where b_n are either 0 or 1.
        
        Args:
            bin_seq: [b_1, b_2, b_3, ...]
        
        Returns:
            decimal_num: The number that corresponds to the input bin_seq
    """
    half_powered = np.array([1/2**k for k in range(1, len(bin_seq)+1)])
    return np.dot(np.array(bin_seq), half_powered)

def arithmetic_decoding(alphabeth, pmf, encoded_signal, num_to_decode):
    """Decode binary signal to a signal with symbols from the alphabeth.
    
    Args:
        alphabeth: List of unique characters.
        pmf: Probability mass function for elements in alphabeth. The order follows
             the order of alphabeth, that is pmf[i] is the probability of alphabeth[i].
        encoded_signal: List of 0 or 1.
        num_to_decode: Number of symbols to decode.
    
    Returns: 
        signal: String with symbols from alphabeth. The decoded signal.
    """
    # Initial step
    decimal_signal = decimal_from_binary(encoded_signal)
    orig_interval_edges = interval_edges(pmf)
    curr_symbol_inds = interval_in_increasing(decimal_signal, orig_interval_edges)
    symbols = [alphabeth[curr_symbol_inds[0]]]
    new_min = orig_interval_edges[curr_symbol_inds[0]]
    new_max = orig_interval_edges[curr_symbol_inds[1]]
    for num_decoded in range(1, num_to_decode):
        curr_interval_edges = scale_interval_edges(orig_interval_edges, new_min, new_max)
        curr_symbol_inds = interval_in_increasing(decimal_signal, curr_interval_edges)
        symbols.append(alphabeth[curr_symbol_inds[0]])
        new_min = curr_interval_edges[curr_symbol_inds[0]]
        new_max = curr_interval_edges[curr_symbol_inds[1]]

    return ''.join(symbols)

alphabeth = ['a', 'b', 'c']
pmf = np.array([0.6, 0.2, 0.2]) # Probability mass function for alphabeth (must sum to 1)
signal = 'abccba'
intervals = arithmetic_intervals(alphabeth, signal, pmf)

for interv in intervals:
    print(interv[0], interv[1])

print(binary_sequence(0.853515625))

decimal_interval = (0.06752, 0.0688)
print(shortest_binary(decimal_interval))

encoded_signal = arithmetic_encoding(alphabeth, pmf, signal)
print(encoded_signal)

encoded_signal = [1, 0, 0, 0, 1]
decoded_signal = arithmetic_decoding(alphabeth, pmf, encoded_signal, 5)
print(decoded_signal)

alphabeth = ['a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z']
pmf = np.array([8.167, 1.492, 2.782, 4.253, 12.702, 2.228, 2.015, 6.094, 6.966, 0.153, 0.772, 4.025, 2.406, 6.749, 7.507, 1.929, 0.095, 5.987, 6.327, 9.056, 2.758, 0.978, 2.360, 0.150, 1.974, 0.074]) / 100
signal = 'helloworld'

encoded_signal = arithmetic_encoding(alphabeth, pmf, signal)
decoded_signal = arithmetic_decoding(alphabeth, pmf, encoded_signal, len(signal))

print('Signal:         {}'.format(signal))
print('Encoded signal: {}'.format(''.join([str(s) for s in encoded_signal])))
print('Decoded signal: {}'.format(decoded_signal))