# This file is part of librefdatool. librefdatool is free software: you can
# redistribute it and/or modify it under the terms of the GNU General Public
# License as published by the Free Software Foundation, version 2.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 51
# Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
# Copyright (C) 2013 Javier D. Garcia-Lasheras
#
# Analysis class code is based in previous open source work:
# * mfreqz: by Matti Pastell
# * zplane: by Chris Felton


import numpy as np
from numpy import pi, log10
from scipy import signal
from matplotlib import pyplot as plt
from matplotlib import mlab
from  matplotlib import patches
from matplotlib.figure import Figure
from matplotlib import rcParams



class Analysis():
    '''This class include tools for LTI filter analysis,
    sporting math models for studying coefficient quantization effects.
    Note (limited to FIR filters)
    '''

    def analyze_zero_pole(self, c, p, q, filename=None):
        '''Plot graphical zero/pole analysis in Z-plane:
        * c: FIR filter coefficients array.
        * filename: optional file for storing plot and not showing it.
        '''
        print('Plotting Z Analysis...')

        # Quantize coefficients
        d = np.zeros(len(c))
        for ii in range(len(c)):
            d[ii] = int(c[ii]*(2**(p+q-1)))

        # Temporal assignation: only valid for FIR 

        print('Coefficients')

        b1 = c
        a1 = 1

        print(b1)
        print(np.max(b1))

        b2 = d
        a2 = 1

        print(b2)
        print(np.max(b1))

        plt.figure('libre-fdatool analysis'); 

        # get a figure/plot
        ax = plt.subplot(111)

        # create the unit circle
        uc = patches.Circle((0,0), radius=1, fill=False,
                        color='black', ls='dashed')
        ax.add_patch(uc)

        # 1 - The coefficients are less than 1, normalize the coeficients
        if np.max(b1) > 1:
            kn1 = np.max(b1)
            b1 = b1/float(kn1)
        else:
            kn1 = 1
    
        if np.max(a1) > 1:
            kd1 = np.max(a1)
            a1 = a1/float(kd1)
        else:
            kd1 = 1
    
        # 2 - The coefficients are less than 1, normalize the coeficients
        if np.max(b2) > 1:
            kn2 = np.max(b2)
            b2 = b2/float(kn2)
        else:
            kn2 = 1
    
        if np.max(a2) > 1:
            kd2 = np.max(a2)
            a2 = a2/float(kd2)
        else:
            kd2 = 1
            
        # Get the poles and zeros
        p1 = np.roots(a1)
        z1 = np.roots(b1)
        k1 = kn1/float(kd1)

        p2 = np.roots(a2)
        z2 = np.roots(b2)
        k2 = kn2/float(kd2)
    
        # 1 - Plot the zeros and set marker properties    
        tz1 = plt.plot(z1.real, z1.imag, 'bo', ms=10)
        plt.setp( tz1, markersize=10.0, markeredgewidth=1.0,
              markeredgecolor='b', markerfacecolor='b')

        # 1 - Plot the poles and set marker properties
        tp1 = plt.plot(p1.real, p1.imag, 'bx', ms=10)
        plt.setp( tp1, markersize=12.0, markeredgewidth=3.0,
                  markeredgecolor='b', markerfacecolor='b')
    
        # 2 - Plot the zeros and set marker properties    
        tz2 = plt.plot(z2.real, z2.imag, 'ro', ms=10)
        plt.setp( tz2, markersize=10.0, markeredgewidth=1.0,
                  markeredgecolor='r', markerfacecolor='r')

        # 2 - Plot the poles and set marker properties
        tp2 = plt.plot(p2.real, p2.imag, 'rx', ms=10)
        plt.setp( tp2, markersize=12.0, markeredgewidth=3.0,
                  markeredgecolor='r', markerfacecolor='r')


        # set axis
        plt.axis('scaled')

        plt.ylabel('Imaginary Component')
        plt.xlabel('Real Component')

        plt.title(r'Zero-Pole Diagram (Blue=Float; Red=Int)')
        plt.grid(True)

        # show or store plot
        return plt
    


    def analyze_frequency_response(self, c, p, q, filename=None):
        '''Plot graphical Magnitude/Phase analysis in frequency domain:
        * c: FIR filter coefficients array.
        * filename: optional file for storing plot and not showing it.
        '''

        # Quantize coefficients
        d = np.zeros(len(c))
        for ii in range(len(c)):
            d[ii] = int(c[ii]*(2**(p+q-1)))
        d = d/(2**(p+q-1))


        print(c)
        print(d)

        wc,hc = signal.freqz(c,1)
        hc_dB = 20 * log10 (abs(hc))

        wd,hd = signal.freqz(d,1)
        hd_dB = 20 * log10 (abs(hd))

        plt.figure('libre-fdatool analysis'); 

        plt.subplot(211)
        plt.plot(wc/max(wc),hc_dB,'b')
        plt.plot(wd/max(wd),hd_dB,'r')

        plt.ylim(-150, 5)
        plt.ylabel('Magnitude (dB)')
        plt.xlabel(r'Normalized Frequency (x$\pi$rad/sample)')
        plt.title(r'Magnitude response (Blue=Float; Red=Int)')
        plt.grid(True)

        plt.subplot(212)
        hc_Phase = np.unwrap(np.arctan2(np.imag(hc),np.real(hc)))
        hd_Phase = np.unwrap(np.arctan2(np.imag(hd),np.real(hd)))
        plt.plot(wc/max(wc),hc_Phase,'b')
        plt.plot(wd/max(wd),hd_Phase,'r')

        plt.ylabel('Phase (radians)')
        plt.xlabel(r'Normalized Frequency (x$\pi$rad/sample)')
        plt.title(r'Phase response (Blue=Float; Red=Int)')
        plt.grid(True)

        plt.subplots_adjust(hspace=0.5)

        return plt