# -*- coding: utf-8 -*-
"""
Created on Mon Apr 11 21:30:24 2016

@author: Yoshi
"""

import numpy as np
from numpy import pi
from scipy import stats
import matplotlib.pyplot as plt
import os

plt.ioff()
np.random.seed(99)

fs = 200                                # Sample frequency (in Hz)
dt = 1. / fs                            # Sample period (s)
maxtime = 2                             # Maximum time in seconds
N = fs * maxtime                        # Total number of samples
t = np.linspace(0, maxtime - dt, N)     # Time array

# Sine wave attributes: frequencies (Hz) and amplitudes
f1 = 6      ;       A1 = 1
f2 = 5      ;       A2 = 1
f3 = 25     ;       A3 = 1

sin1 = A1*np.sin(2*pi*f1*t)
sin2 = A2*np.sin(2*pi*f2*t)
sin3 = A3*np.sin(2*pi*f3*t)

sin_total = sin1

# Code from Notes
A = sin_total                           # Input signal
B = np.fft.fft(A)                       # FFT of signal
C = np.zeros(N, dtype='complex128')     # Output array initalization

# Hilbert Transform
C[0] = B[0]                     # Hard-code start and
C[N/2] = B[N/2]                 # Nyquist of data output

for i in range(1, N/2):         # Double positive frequencies
    C[i] = 2.0 * B[i]

for i in range(N/2 + 1, N):     # Zero negative frequencies
    C[i] = 0.0

B = np.fft.ifft(C)              # Inverse FFT of C: Store in B

# Plot Function and Envelope Function
fig = plt.figure()
fig.text(0.1, -0.01, r'$f_1 = %dHz$, $f_2 = %dHz$' % (f1, f2), fontsize=16)

ax = fig.add_subplot(1,1,1)

ax.plot(t, A, 'b')
ax.plot(t, abs(B), 'r', lw=2)
ax.set_title('Signal Envelope of Two Closely Spaced Sinusoids')
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude')

# Look at Phase Information
phase = np.angle(B)              # Used because np.arctan won't work with ASTC
unwrapped = np.unwrap(phase)     # Write own version of this