import unittest, submitted
from gradescope_utils.autograder_utils.decorators import weight
import numpy as np
import librosa

# TestSequence
class TestStep(unittest.TestCase):
    @weight(10.0)
    def test_welch_199frames(self):
        signal = np.arange(10000)
        frames, pgram, spectrum = submitted.welch(signal, np.bartlett(100))
        self.assertEqual(len(frames.shape), 2, 'frames should be an (nframes,framelength) array')
        self.assertEqual(len(pgram.shape), 2, 'pgram should be an (nframes,framelength) array')
        self.assertEqual(len(spectrum.shape), 1, 'spectrum should be a 1d array')
        self.assertEqual(frames.shape[0], 199, 'frames should have 199 frames')
        self.assertEqual(pgram.shape[0], 199, 'pgram should have 199 frames')
        self.assertEqual(frames.shape[1], 100, 'framelength should be 100')
        self.assertEqual(pgram.shape[1], 100, 'pgram framelength should be 100')
        self.assertEqual(max(spectrum.shape), 100, 'spectrum length should be 100')
        self.assertAlmostEqual(frames[0,45], signal[45], places=1,
                               msg='frames[0,45] should equal signal[45]')
        self.assertAlmostEqual(frames[12,45], signal[12*50+45], places=1,
                               msg='frames[m,n] should equal signal[m*L/2+n]')
        
    @weight(10.0)
    def test_welch_1frame(self):
        signal = np.arange(85)
        frames, pgram, spectrum = submitted.welch(signal, np.hanning(80))
        self.assertEqual(max(frames.shape), 80, 'frames should have length 80')
        self.assertEqual(max(pgram.shape), 80, 'frames should have length 80')
        self.assertEqual(max(spectrum.shape), 80, 'spectrum length should be 80')
        
    @weight(10.0)
    def test_lpc_1Hz(self):
        sample_rate = 1
        f = np.random.uniform(low=0.0001, high=0.4999, size=5)*sample_rate
        omega = f * 2*np.pi/sample_rate
        b = np.random.uniform(low=0.0001, high=1.9999, size=5)*sample_rate
        sigma = b * np.pi/sample_rate
        p = np.concatenate((np.exp(-sigma+1j*omega), np.exp(-sigma-1j*omega)))
        a = np.poly(p)
        frequency, bandwidth = submitted.lpc_roots(a, sample_rate)
        self.assertEqual(max(frequency.shape), 10, 'there should be 10 frequencies')
        self.assertEqual(max(bandwidth.shape), 10, 'there should be 10 bandwidths')
        self.assertAlmostEqual(np.amin(np.abs(frequency-f[3])), 0, places=1,
                               msg='frequency should contain %g at Fs=%g'%(f[3],sample_rate))
        self.assertAlmostEqual(np.amin(np.abs(bandwidth-b[3])), 0, places=1,
                               msg='bandwidth should contain %g at Fs=%g'%(b[3],sample_rate))
        
    @weight(10.0)
    def test_lpc_100Hz(self):
        sample_rate = 100
        f = np.random.uniform(low=0.0001, high=0.4999, size=4)*sample_rate
        omega = f * 2*np.pi/sample_rate
        b = np.random.uniform(low=0.0001, high=1.9999, size=4)*sample_rate
        sigma = b * np.pi/sample_rate
        p = np.concatenate((np.exp(-sigma+1j*omega), np.exp(-sigma-1j*omega)))
        a = np.poly(p)
        frequency, bandwidth = submitted.lpc_roots(a, sample_rate)
        self.assertEqual(max(frequency.shape), 8, 'there should be 8 frequencies')
        self.assertEqual(max(bandwidth.shape), 8, 'there should be 8 bandwidths')
        self.assertAlmostEqual(min(np.abs(frequency-f[3])), 0, places=1,
                               msg='frequency should contain %g at Fs=%g'%(f[3],sample_rate))
        self.assertAlmostEqual(min(np.abs(bandwidth-b[3])), 0, places=1,
                               msg='bandwidth should contain %g at Fs=%g'%(b[3],sample_rate))
        
    @weight(10.0)
    def test_rls(self):
        memory_weight = 0.95
        P_init = np.eye(10)
        signal = np.arange(1000)*np.power(-1,np.arange(1000))
        predictor = submitted.rls(signal, P_init, memory_weight)
        self.assertEqual(len(predictor.shape),2,'predictor should be a 2-array')
        self.assertEqual(predictor.shape[0],990,'predictor should have 990 rows')
        self.assertEqual(predictor.shape[1],10,'predictor should have 10 columns')
        self.assertLess(np.sum(predictor),1,'sum(predictor) should be less than -1')
        
    @weight(10.0)
    def test_rls(self):
        memory_weight = 0.99
        P_init = np.eye(5)
        signal = np.arange(800)
        predictor = submitted.rls(signal, P_init, memory_weight)
        self.assertEqual(len(predictor.shape),2,'predictor should be a 2-array')
        self.assertEqual(predictor.shape[0],795,'predictor should have 795 rows')
        self.assertEqual(predictor.shape[1],5,'predictor should have 5 columns')
        self.assertGreater(np.sum(predictor),1,'sum(predictor) should be greater than 1')
        
