"""
Model Butterworth filter, including effects of bit truncation.  There
are a couple of sign things I don't understand exactly but this is a
good start.

MFH 2014-07-15
"""

import numpy as np

def bit_mask(n):
    return 2**n - 1

def promote(x,n1,n2):
    # Sign extend the n1-bit signed integer in x to be an n2-bit
    # signed integer.
    x = x & bit_mask(n1)
    if x & (2**(n1-1)):
        x |= bit_mask(n2-n1) << n1
    return x

def resign(x, n, str=False):
    # Sign extend from bit n and return signed integer.  Maybe string.
    if str:
        nz = (n+3)/4
        fmt = '%%+0%ix' % (nz+1)
        return fmt % resign(x, n)
    x = x & bit_mask(n)
    if x & (2**(n-1)):
        return int(x - 2**n)
    return int(x)

def sbmul(a, na, b, nb=None):
    # Multiply a (na bits signed) by b (which is considered to be
    # signed if nb is specified.
    if not nb is None:
        b = resign(b, nb)
    return resign(a, na) * b


default_params = {
    # Number of significant bits in the input word
    'input_width': 29,
    # Width in bits of the WN registers:
    'accumulator_width': 29,
    # Number of significant bits in the B coefficients.
    'coefficient_width': 15,
    # Windowing applied before adding X to (W1*B1+W2*B2)>>(coeff_width-1)
    'window_width': 29,
    # Number of bits to away in the output
    'output_shift': 14,
    # Number of bits to keep in the output.
    'output_window': 32,
}

class TButterworth:
    """
    Single biquad Butterworth filter.
    """

    def __init__(self, (B1, B2), params):
        self.wn = np.zeros(2, 'int64')
        self.b = np.zeros(2, 'int64')
        self.verbose = False
        self.params = default_params.copy()
        self.params.update(params)
        self.b[:] = B1, B2

    def vprint(self, n):
        if self.verbose:
            print n

    def reset(self):
        self.wn[:] = 0

    def internal_gain(self):
        BBITS = self.params['coefficient_width']
        return 1./ (1 - (self.b[0]-self.b[1])/2.**(BBITS-1))

    def add(self, x):
        WBITS = self.params['accumulator_width']
        BBITS = self.params['coefficient_width']
        ABITS = self.params['window_width']
        self.vprint('ADD ' + resign(x, WBITS, True))
        AB_BITS = WBITS+BBITS
        a = sbmul(self.wn[0], WBITS, self.b[1]) & bit_mask(AB_BITS)
        b = sbmul(self.wn[1], WBITS, self.b[0]) & bit_mask(AB_BITS)
        self.vprint(' a,b = %s,%s' % 
                    (resign(a, AB_BITS, True), resign(b, AB_BITS, True)))
        c = ((b-a) >> (BBITS - 1)) & bit_mask(ABITS)
        self.vprint(' c = a-b = %s' % resign(c, ABITS))
        IBITS = self.params['input_width']
        wn = (resign(x, IBITS) + c) & bit_mask(ABITS)
        self.vprint(' X + c = %s' % resign(wn, WBITS))
        out1 = (resign(self.wn[0], WBITS) + 
                resign(self.wn[1], WBITS)*2 + 
                resign(wn, WBITS))
        out2 = out1 >> self.params['output_shift']
        OBITS = self.params['output_window']
        out3 = out2 & bit_mask(OBITS)
        # Store.
        self.wn[0] = self.wn[1] & bit_mask(WBITS)
        self.wn[1] = wn & bit_mask(WBITS)
        self.vprint(' saved = %s %s' % 
                    (resign(self.wn[0], WBITS, True),
                     resign(self.wn[1], WBITS, True)))
        self.vprint(' output = ' + resign(out3, OBITS, True))
        return out3


print 'Steady state test'
for name, width in [('default', 29), ('expanded', 32)]:
    tbut = TButterworth((32568, 16188), {
            'accumulator_width': width,
            'window_width': width,
            })
    print '  ' + name
    g = tbut.internal_gain()
    for test_val in 2**np.arange(24)-1:
        tbut.wn[:] = test_val * g
        for i in range(1000):
            x = tbut.add(test_val)
        y = tbut.add(test_val)
        ok = (x == tbut.add(test_val))
        print '   %12i -> %12i %12i  %s' % (test_val, x, y, ok)
    print