# -*- coding: utf-8 -*-
"""
Created on Mon Sep 28 13:18:36 2020

@author: hmharri3
XXX.py p

Input defined in processing options and plot options sections.



<333 Trey 09/2020
"""

# import packages

import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
import pyswarms as ps

# define signal to reconstruct

fs        = 3600
fLo       = 0
fHi       = fs/2
N         = 3600
# fKnown = np.asarray([10])
# XKnown = np.asarray([10])
dcOffset    = 0
fKnown = np.asarray([8,48,72,80,88,96,104])
XKnown   = np.asarray([0.6229356,0.085519664,0.18917666,0.4036676,
                           0.15955658,2.217998,0.49061242,])

# construct frequency content of signal

k  = np.arange(N)
fk = k*fs/N
X  = np.zeros(len(fk))

for i,f in enumerate(fk):
    if f in fKnown:
        X[i] = XKnown[f == fKnown]

plt.figure()
plt.plot(fk,X)
plt.xlim(0,120)

# calculate inverse discrete fourier transform

dt = 1/fs * 360
x = N*np.real(np.fft.ifft(X))
x = (x - np.mean(x))
t = np.arange(N)*dt


# convert to degrees and only take one revolution
theta = np.append(t[t<360],360)
trueSignal  = np.append(x[t<360],x[0]) + 9



# solve for coefficients

# change variable names to make them spatial instead of temporal
waveNumber = np.atleast_1d([8,96,104])

# build design matrix
# probeTheta = np.round(np.random.rand(8)*360,1)
def solveForWaveletCoefficients(probeTheta, waveNumber, theta, trueSignal):
    m          = len(probeTheta)
    N          = len(waveNumber)*2 + 1
    A          = np.zeros([m,N])
    x          = np.zeros([m])
    tsInterp = interp1d(theta, trueSignal, kind='linear')
    for i,pt in enumerate(probeTheta):
        x[i] = tsInterp(pt)
        for j in np.arange((N-1)//2):
            A[i,2*j]   = np.sin(waveNumber[j]*pt*np.pi/180)
            A[i,2*j+1] = np.cos(waveNumber[j]*pt*np.pi/180)
    # last column equals 1
    A[:,-1] = 1
    F = np.linalg.lstsq(A,x,rcond=None)[0]
    return A, F

# define objective function

# goal is to minimize condition number of design matrix
# low condition numbers give lowest error metrics

def cost_function(probeTheta, waveNumber, A):
    if len(probeTheta) == len(waveNumber)*2+1:
        k = np.linalg.norm(A)*np.linalg.norm(np.linalg.inv(A))
    else:
        k = np.linalg.norm(A)*np.linalg.norm(np.linalg.pinv(A))
    fcost = np.sum((1*k)**2)
    return fcost

# define optimization function

def opt_func(probeTheta):
    # waveNumber, theta, trueSignal = getOptParameters()
    n_particles = probeTheta.shape[0]
    cost = np.zeros(n_particles)
    for i in range(n_particles):
        A, F = solveForWaveletCoefficients(probeTheta[i], waveNumber, theta, trueSignal)
        cost[i] = cost_function(probeTheta[i], waveNumber, A)
    # cost = [cost_function(probeTheta[i],)]
    return cost

# define swarm parameters

swarm_size = 500
dim = 9 # Dimension of X
epsilon = 1.0
options = {'c1': 1.5, 'c2':1.5, 'w':0.5}
probeSpacing = 360/dim
tLo     = np.zeros(dim)
tHi     = np.zeros(dim)
minSpacing = 3
for i in range(dim):
    tLo[i] = probeSpacing*i
    if i != dim-1:
        tHi[i] = probeSpacing*(i+1) + minSpacing
    else:
        tHi[-1] = probeSpacing*(i+1) - minSpacing
constraints = (tLo,tHi)
(np.ones(dim)*0, np.ones(dim)*360)

#  do the optimization

# call optimizer instance of pso
optimizer = ps.single.GlobalBestPSO(n_particles=swarm_size, dimensions=dim, 
                                    options=options, bounds=constraints)

# optimize
cost, joint_vars = optimizer.optimize(opt_func, iters=200)

from pyswarms.utils.plotters import plot_cost_history
plt.figure()
plot_cost_history(optimizer.cost_history)

# reconstruct signal

A, F = solveForWaveletCoefficients(joint_vars, waveNumber, theta, trueSignal)
N          = len(waveNumber)*2 + 1
xR = F[-1]
xR = np.zeros(len(theta))
for i,t in enumerate(theta):
    xRtemp = 0
    for j in np.arange((N-1)//2):
        xRtemp = xRtemp + (F[2*j]*np.sin(waveNumber[j]*t*np.pi/180) +
                          F[2*j+1]*np.cos(waveNumber[j]*t*np.pi/180))
    xR[i] = xRtemp + F[-1]

plt.figure()
plt.plot(theta,trueSignal)
plt.plot(theta,xR,'--')
plt.xlim((0,60))

