# -*- coding: utf-8 -*-
"""
Created on Sat Apr 09 12:10:33 2016

@author: Jeremy
"""
import numpy as np
import scipy.signal
import matplotlib.pyplot as plt


def load_flight_data(filename):
    """
    Returns the data from flight 18 provided for HW1, as a dictionary.
    """
    return np.genfromtxt(filename, deletechars='%', skip_header=2,
                         names=('time', 'u', 'v', 'w', 'T', 'q', 'p'))


def get_power_spectral_density(timeseries, rate=25., nfft=4096, noverlap=None):
    """
    Returns the power spectral density of the given timeseries.

    Parameters
    ----------
    timeseries : ndarray
        One-dimensional timeseries for which we want the PSD.
    rate : float, optional
        Sampling rate of the data, in Hz.
    nfft : int, optional
        Window size for sampling power spectral density, in number of samples.
    noverlap : int, optional
        Number of points to overlap between segments. Defaults to nfft/2.

    Returns
    -------
    f : ndarray
        The frequency axis corresponding to Pxx.
    Pxx : ndarray
        The power spectral density of the given timeseries.
    """
    f, Pxx = scipy.signal.welch(
        timeseries,
        fs=rate,
        nperseg=nfft,
        noverlap=noverlap,
        window=np.hanning(nfft),
        detrend=lambda x: x,  # no detrending. detrend=False should also work
        scaling='density')
    return f, Pxx


def apply_high_pass_filter(timeseries, rate=25., cutoff_frequency=0.1):
    """
    Takes in a timeseries, and returns a timeseries with a 4th-order
    butterworth high-pass filter of the given cutoff frequency (in Hz) applied.
    Rate is the frequency of the incoming timeseries, in Hz.
    """
    b, a = scipy.signal.butter(N=4, Wn=cutoff_frequency/rate, btype='high')
    high_passed = scipy.signal.filtfilt(b=b, a=a, x=timeseries)
    initialization_error_cutoff = round(rate/cutoff_frequency)
    return high_passed[initialization_error_cutoff:]


if __name__ == '__main__':
    data_filename = 'rf18L1.txt'
    data = load_flight_data(data_filename)

    # Question 2
    f, Pww = get_power_spectral_density(data['w'])

    # Question 4
    whi = apply_high_pass_filter(data['w'])