#!/usr/bin/python
# -*- coding: utf8 -*-

import numpy as np
import scipy.signal as sig
import matplotlib as mpl
import matplotlib.pyplot as plt
from math import *

N = 5 # order of the filters
band = 1.0 # cut-off normalized frequency
Rpass = 1.0 # ripple in the pass-band (dB)
Rstop = 20 # ripple in the stop-band (dB)
frange = np.linspace(0, 2, 2001)
color = 'r'
fname = 'filters_order5.svg'

mpl.rcParams['axes.grid'] = True
plt.figure(figsize=(6, 5.4))

(num, den) = sig.butter(N, band, analog=True)
filterfy = sig.freqs(num, den, frange)
plt.subplot(221)
plt.title('Butterworth')
plt.plot(filterfy[0], np.abs(filterfy[1]), lw=3, color=color)

(num, den) = sig.cheby1(N, 1, band, analog=True)
filterfy = sig.freqs(num, den, frange)
plt.subplot(222)
plt.title('Chebyshev type 1')
plt.plot(filterfy[0], np.abs(filterfy[1]), lw=3, color=color)

(num, den) = sig.cheby2(N, Rstop, band, analog=True)
filterfy = sig.freqs(num, den, frange)
plt.subplot(223)
plt.title('Chebyshev type 2')
plt.plot(filterfy[0], np.abs(filterfy[1]), lw=3, color=color)

(num, den) = sig.ellip(N, Rpass, Rstop, band, analog=True)
filterfy = sig.freqs(num, den, frange)
plt.subplot(224)
plt.title('Elliptic')
plt.plot(filterfy[0], np.abs(filterfy[1]), lw=3, color=color)

for ax in plt.gcf().get_axes():
    ax.set_xlim(*frange[[0,-1]])
    ax.set_ylim(0, 1.1)
    ax.set_xlabel('f/f$_0$')
    ax.xaxis.set_label_coords(0.65, -0.1)
    ax.set_ylabel('G', rotation=0)
    ax.yaxis.set_label_coords(-0.05, 1.0)

plt.tight_layout(pad=1, w_pad=2, h_pad=2)
plt.savefig(fname)

def postprocess(fname):
    from lxml import etree
    with open(fname, 'r') as svgfile:
        tree = etree.parse(svgfile, etree.XMLParser(remove_blank_text=True))
    svg = tree.getroot()
    nsmap = '{' + svg.nsmap[None] + '}'

    # move all definitions to the front
    svg[:] = sorted(svg, key=lambda el: {False:0, True:1}[el.tag!=nsmap+'defs'])
    with open(fname, 'w') as svgfile:
        tree.write(svgfile,
            xml_declaration=True, pretty_print=True, encoding='utf-8')

postprocess(fname)