#!/usr/bin/env python

import numpy
import numpy.linalg
import scipy
import scipy.integrate
import scipy.linalg

def eigenfilter(M, P, v):
    U = scipy.linalg.null_space(M)
    eigs = numpy.linalg.eigh(U.T @ P @ U)
    assert eigs[0][0] > 0, "insufficient constraints for good conditioning"
    return (U @ eigs[1][:,0]) / (v @ U @ eigs[1][:,0])

def fourier(n, freq, center=0.5):
    # center = 0    ==>  first sample is index 0
    # center = 0.5  ==>  middle sample is index 0
    index = numpy.arange(n) - (n - 1) * center
    return numpy.exp(2j * numpy.pi * freq * index)

def fourier_integral(n, freq0, freq1, center=0.5):
    index = numpy.arange(n) - (n - 1) * center
    def integral(freq):
        return (numpy.where(index == 0, freq, fourier(n, freq, center)) /
                numpy.where(index == 0, 1, 2j * numpy.pi * index))
    return integral(freq1) - integral(freq0)

def linear_phase(n, even=True):
    return numpy.eye(n) + numpy.rot90(numpy.eye(n)) * (-1 if even else 1)

def band_energy(n, freq0, freq1):
    assert 0 <= freq0 <= freq1 <= 0.5
    return scipy.linalg.toeplitz(2 * freq1 * numpy.sinc(2 * freq1 * numpy.arange(n)) -
                                 2 * freq0 * numpy.sinc(2 * freq0 * numpy.arange(n)))

def band_deviation(n, freq0, freq1, reference, center=0.5):
    assert 0 <= freq0 <= freq1 <= 0.5
    def response_error(freq):
        error_vector = fourier(n, freq, center) - reference
        return 2 * numpy.real(numpy.outer(numpy.conj(error_vector), error_vector))
    return scipy.integrate.quad_vec(response_error, freq0, freq1)[0]

def band_deviation_lowpass(n, freq, center=0.5):
    assert 0 <= freq <= 0.5
    x = numpy.arange(n) - (n - 1) * center
    y = numpy.arange(n).reshape((-1, 1)) - (n - 1) * center
    return 2 * freq * (1 + numpy.sinc(2 * freq * (x - y))
                       - numpy.sinc(2 * freq * x) - numpy.sinc(2 * freq * y))

def arbitrary_freq_response(n, reference, weight, center=0.5):
    def response_error(freq):
        error_vector = fourier(n, freq, center) - reference(freq)
        outer = numpy.outer(numpy.conj(error_vector), error_vector)
        return weight(freq) * numpy.real(outer)
    return 2 * scipy.integrate.quad_vec(response_error, 0, 0.5)[0]

def time_delay(n):
    return numpy.diag(numpy.arange(n)**2)

def linear_constraint(A, b, v):
    error_matrix = A - numpy.outer(b, v)
    return numpy.real(numpy.conj(error_matrix.T) @ error_matrix)

def filter_plot(filters, filename=None, center=0.5):
    import pylab
    fig, (a0, a1) = pylab.subplots(1, 2, figsize=(8, 3))
    xmin, xmax, rmax = 0, 0, 0
    for name, h in filters:
        n = len(h)
        x = numpy.arange(n) - (n - 1) * center
        xmin = min(xmin, min(x) - 0.5)
        xmax = max(xmax, max(x) + 0.5)
        a0.plot(x, h, '.-', label=name, lw=.2)
        resp = numpy.fft.rfft(h, n=2000)
        rmax = max(rmax, max(abs(resp)))
        f = numpy.linspace(0, .5, len(resp))
        a1.plot(f, numpy.abs(resp), lw=2, label=name)
    a0.set(title="Impulse response", xlim=(xmin, xmax),
           xlabel="tap number", ylabel="weight")
    a1.set(title="Frequency response", xlim=(0, 0.5), ylim=(0, rmax),
           xlabel="frequency / sample rate", ylabel="magnitude")
    a0.grid()
    a1.grid()
    if len(filters) > 1:
        a1.legend()
    pylab.tight_layout()
    if filename:
        pylab.savefig(filename)

def lowpass(n, passband, stopband):
    return eigenfilter(linear_phase(n),
                       band_deviation_lowpass(n, passband) + band_energy(n, stopband, 0.5),
                       numpy.ones(n))

filter_plot([("{:.2f}, {:.2f}".format(i, j), lowpass(31, i, j))
             for i, j in [(0, 0.1), (.1, .2), (.05, .1)]],
            "lowpass.svg")

def lowpass_minimum_time(n, stopband, time_weight): 
    return eigenfilter(numpy.zeros((n, n)),
                       band_energy(n, stopband, 0.5) + time_delay(n) * time_weight,
                       numpy.ones(n))

filter_plot([("{:.2f}, 1e{}".format(i, j), lowpass_minimum_time(31, i, 10**j))
             for i, j in [(0.1, -4), (0.2, -6), (0.3, -8)]],
            "lowpass_minimum_time.svg", center=0)

def bandpass(n, passbands, stopbands):
    v = numpy.real(sum([fourier_integral(n, p[0], p[1]) for p in passbands])
                   / sum([p[1] - p[0] for p in passbands]))
    P = (sum([band_deviation(n, b[0], b[1], v) for b in passbands])
         + sum([band_energy(n, b[0], b[1]) for b in stopbands]))
    return eigenfilter(linear_phase(n), P, v)

passbands = [(.1, .2), (.35, .4)]
stopbands = [(0, .05), (.25, .3), (.45, .5)]
filter_plot([("{}".format(i), bandpass(i, passbands, stopbands)) for i in [63, 31, 15]],
            "bandpass.svg")

def hilbert(n, f_min, f_max):
    v = numpy.imag(fourier_integral(n, f_min, f_max)) / (f_max - f_min)
    return eigenfilter(linear_phase(n, False), band_deviation(n, f_min, f_max, 1j * v), v)

filter_plot([("{}, {:.2f}, {:.2f}".format(n, i, j), hilbert(n, i, j))
             for n, i, j in [(15, .1, .4), (16, .05, .5)]], "hilbert.svg")

def differentiator(n, passband, stopband):
    v = numpy.arange(n) - (n - 1) / 2
    desired_response = lambda freq: 2j * numpy.pi * freq * v
    frequency_weight = lambda freq: 1 if freq < passband else 0
    P = (arbitrary_freq_response(n, desired_response, frequency_weight)
         + band_energy(n, stopband, 0.5))
    return eigenfilter(linear_phase(n, False), P, v)

filter_plot([("{:.2f}, {:.2f}".format(i, j), differentiator(31, i, j))
             for i, j in [(0, .1), (.05, .10), (.1, .2)]],
            "differentiator.svg")

def half_sample_delay(n, passband, stopband, center=0.5):
    desired_response = lambda freq: numpy.exp(1j * numpy.pi * freq)
    frequency_weight = lambda freq: 0.5 / freq if freq < passband else 0
    P = (arbitrary_freq_response(n, desired_response, frequency_weight, center)
         + band_energy(n, stopband, 0.5))
    return eigenfilter(numpy.zeros((n, n)), P, numpy.ones(n))

filter_plot([("{:.2f}, {:.2f}".format(i, j), half_sample_delay(19, i, j))
             for i, j in [(.3, .45), (.4, .5), (.5, .5)]],
            "half_sample_delay.svg")

filter_plot([("{:.2f}, {:.2f}".format(i, j), half_sample_delay(19, i, j, 0))
             for i, j in [(.3, .35), (.4, .43), (.5, .5)]],
            "half_sample_delay_causal.svg", center=0)