import numpy as np
import matplotlib.pyplot as plt
import scipy.linalg


def model_func(p, x):
    func = np.exp(-(x / p[1]) + 1j * 2 * np.pi * p[0] * x)

    return func


def time_evol(p, x):
    return p[0] * np.exp(-(x / p[1]) ** p[2])


def create_data(m, p, n, x1, x2):
    data_matrix = np.zeros([len(x1), len(x2)], dtype=complex)
    ideal = np.zeros_like(x2)
    for k, v in enumerate(m):
        _x1 = v[0](p[k][0], x1)
        _x2 = v[1](p[k][1], x2)
        ideal += v[1](p[k][1], x2)
        data_matrix += np.outer(_x1, _x2)
    data_matrix.real += noise * np.amax(data_matrix).real * np.random.randn(len(x1), len(x2))
    data_matrix.imag += noise * np.amax(data_matrix).real * np.random.randn(len(x1), len(x2))
    return data_matrix, ideal


# ------------------------------------------------------------------------------------------------ #

sw = 100e3  # sampling rate
t = np.arange(1 * 1024) / sw  # time scale
tau = np.logspace(-5, 3, 27)
no_measurements = len(tau)
no_pts = len(t)

f = np.fft.fftshift(np.fft.fftfreq(no_pts, d=1 / sw))

noise = 0.01
models = [[model_func, time_evol],
          [model_func, time_evol],
          # [model_func, time_evol],
          ]
p = [[[-10e3, 5e-4], [0.2, 0.1, 1]],
     [[10e3, 5e-4], [1, 10, 1]],
     # [[0,  5e-4], [1, 1, 1]],
     ]

real_data, ideal_curve = create_data(models, p, noise, t, tau)
# real_data = np.fft.fftshift(np.fft.fft(real_data, axis=0), axes=(0,))

if True:
    f1, axarr1 = plt.subplots()
    for i in range(no_measurements):
        axarr1.plot(f, real_data[:, i].real / np.max(real_data.real) + i, 'b-')
        axarr1.plot(f, real_data[:, i].imag / np.amax(real_data.real) + i, 'r-')

# raw_data = np.append(real_data.real, real_data.imag, axis=0).T
raw_data = real_data.T
spec_sum = np.sum(raw_data.real, axis=1)

# ----------------------- PCA -------------------------------
# calculate deviation
meanvalues = raw_data.mean(axis=0)
data_no_mean = raw_data - meanvalues
# data_scaled = data_no_mean/data_no_mean.std(axis=0)

data = data_no_mean

# SVD
# u_full, s_full, vt_full = np.linalg.svd(data, full_matrices=False)
s_full, u_full = np.linalg.eig(np.cov(data))
vt_full = u_full.T
# ss_full = np.diag(s_full)
eigenvalues = s_full / sum(s_full)

fig1, (ax1, ax2) = plt.subplots(2)

sum_eigen = 0
n = 0
eigen_i = np.inf
while sum_eigen < 98 and n < 7:
    eigen_i = 100 * eigenvalues[n]
    sum_eigen += eigen_i
    ax1.plot(vt_full[n, :] / max(abs((vt_full[n, :]))) - 2 * n, label='{:.2f}'.format(eigen_i))
    ax2.plot(u_full[:, n], label='{:.2f}'.format(eigen_i))
    n += 1
ax1.legend()
plt.show()

n = int(input('Number of components? '))
print('Use first {} components'.format(n))

# reduce dimensionality
u_cut = u_full[:, :n]
s_cut = np.diag(s_full[:n])
vt_cut = vt_full[:n, :]

v_cut = vt_cut.T
transform_matrix = np.dot(u_cut, s_cut)

# reconstruct data
cut_data = np.dot(transform_matrix, vt_cut)
cut_data += meanvalues

if n == 2:
    scores = np.dot(cut_data, v_cut[:, 0] + v_cut[:, 1])
    # scores = np.dot(cut_data, v_cut[:, 0] / sum(v_cut[:, 0]) + v_cut[:, 1] / sum(v_cut[:, 0]))
else:
    scores = np.dot(cut_data, v_cut[:, 0])
# scores = np.sum(cut_data+_meanvalues, axis=1)
# scores = np.dot(transform_matrix, transform_matrix.T)[:, 0]
# scores2 = np.sum(cut_data[:, :len(cut_data)], axis=1)
scores2 = np.sum(cut_data[:, :5], axis=1)
# scores2 = np.dot(cut_data, cut_data.T)[:, 0]


if True:
    fig, ax = plt.subplots(3)
    for i in range(no_measurements):
        # print i, sum(u_full[i, :]), u_cut[i, :]
        ax[0].plot(raw_data[i, :])
        ax[1].plot(cut_data[i, :])
        ax[2].plot(cut_data[i, :] - raw_data[i, :])
    ax[0].set_title('full')
    ax[1].set_title('PCA')
    ax[2].set_title('difference')

# coeff, score, latent = princomp(data, numpc=2)
# x_r = np.dot(coeff, score).T + meanvalues
# plt.plot(np.sum(score, axis=0))
# plt.show()

if False:
    if n == 2:
        f7, (ax7, ax8, ax9) = plt.subplots(3)
        ax7.plot(u_cut[:, 0], u_cut[:, 1], '-')
        ax8.plot(v_cut[:, 0], v_cut[:, 1], '-')
        ax9.plot(cut_data[:, 0] / sum(v_cut[:, 1]), cut_data[:, 1] / sum(v_cut[:, 0]), '-')
        ax9.plot(cut_data[:, 0], cut_data[:, 1], '-')
        ax7.set_xlabel('C1')
        ax7.set_ylabel('C2')
        ax8.set_xlabel('C1')
        ax8.set_ylabel('C2')
        ax9.set_xlabel('C1')
        ax9.set_ylabel('C2')
    elif n == 3:
        from mpl_toolkits.mplot3d import Axes3D

        f7 = plt.figure()
        ax7 = f7.add_subplot(311, projection='3d')
        ax8 = f7.add_subplot(312, projection='3d')
        ax9 = f7.add_subplot(313, projection='3d')
        ax7.plot(u_cut[:, 0], u_cut[:, 1], u_cut[:, 2])
        ax8.plot(v_cut[:, 0], v_cut[:, 1], v_cut[:, 2])
        # ax9.scatter(cut_data[:, 0]/sum(v_cut[:, 1]), cut_data[:, 1]/sum(v_cut[:, 0]), '-o')
        ax9.plot(cut_data[:, 0], cut_data[:, 1], cut_data[:, 2])
        ax7.set_xlabel('C1')
        ax7.set_ylabel('C2')
        ax7.set_zlabel('C3')
        ax8.set_xlabel('C1')
        ax8.set_ylabel('C2')
        ax8.set_zlabel('C3')
        ax9.set_xlabel('C1')
        ax9.set_ylabel('C2')
        ax9.set_zlabel('C3')

f5, ax5 = plt.subplots(1)
ax5.semilogx(tau, ideal_curve / max(ideal_curve), 'g--', label='model')
ax5.semilogx(tau, np.sum(real_data[:5, :], axis=0).real, label='raw')
ax5.semilogx(tau, spec_sum / max(spec_sum), 'm-', label='spectrum: Re')
ax5.semilogx(tau, scores2 / max(abs(scores2)), 'o', label='PCA')
ax5.semilogx(tau, scores / max(abs(scores)), '^', label='PCA_2')
# ax5.semilogx(tau, (scores - scores[-1]) / (scores[0] - scores[-1]), 'yo', label='PCA_Oleg')
plt.legend()

plt.show()


# #--- try MSSA ---#
# m = 5
# # f6, ax6 = plt.subplots()
# # ax6.plot(real_data)
# y = np.array([], dtype=complex)
# n = len(real_data[:, 0])
# for d in real_data.T:
#     y_i = np.zeros((n, m), dtype=complex)
#     for i in xrange(m):
#         y_i[:n-i, i] = d[i:n]
#     try:
#         y = np.hstack((y, y_i))
#     except ValueError:
#         y = y_i
#
# c = np.dot(y.T, y) / n
# lamb, rho = eig(c)
# pc = np.dot(y, rho)
# for i, k in enumerate(pc.T):
#     plt.plot(k/max(k) + i)
# plt.show()
# rc = []
# for k in xrange(len(real_data[0])):
#     rc_k = np.zeros((n, m), dtype=complex)
#     for i in xrange(m):
#         z = np.zeros((n, m), dtype=complex)
#         for j in xrange(m):
#             z[j:n, j] = pc[:n-j, i]
#         rc_k[:, i] = np.dot(z, rho[k*m:(k+1)*m, i])/m
#         # print rc_k[:, 0].shape
#         # plt.plot(rc_k[:, 0])
#         # plt.show()
#     rc.append(rc_k)
# # fig, ax = plt.subplots(1)
# # for k, v in enumerate(rc):
# #     plt.plot(np.sum(v[:, :2], axis=1))
# #     plt.plot(data_scaled[:, k].real)
# # plt.show()
#
# s_orig = []
# s_ssa = []
# for k, v in enumerate(raw_data.T):
#     s_orig.append(max(v.real)+meanvalues[k])
#     s_ssa.append(max((np.sum(rc[k][:, :], axis=1)*data_std[k]).real)+meanvalues[k])
# # print s_orig
# s_orig = np.array(s_orig)
# s_ssa = np.array(s_ssa)
# plt.plot(s_orig/s_orig[0], label='Orig')
# plt.plot(s_ssa/s_ssa[0], label='SSA')
# # plt.plot(ideal_curve)
# # plt.plot(np.array(s_orig)/np.array(s_ssa))
# plt.legend()
# plt.show()

# n = int(raw_input('Number of components? '))
# print 'Use first {} components'.format(n)
#
# # reduce dimensionality
# u_cut = u_full[:, :n]
# s_cut = np.diag(s_full[:n])
# vt_cut = vt_full[:n, :]
#
# v_cut = vt_cut.T
# transform_matrix = np.dot(u_cut, s_cut)
#
# # reconstruct data
# cut_data = np.dot(transform_matrix, vt_cut)
# cut_data += meanvalues
#
# if n == 2:
#     scores = np.dot(cut_data, v_cut[:, 0] + v_cut[:, 1])
#     scores = np.dot(cut_data, v_cut[:, 0]/sum(v_cut[:, 0]) + v_cut[:, 1]/sum(v_cut[:, 0]))
# else:
#     scores = np.dot(cut_data, v_cut[:, 0])
# # scores = np.sum(cut_data+_meanvalues, axis=1)
# # scores = np.dot(transform_matrix, transform_matrix.T)[:, 0]
# scores2 = np.sum(cut_data[:, :len(cut_data)], axis=1)
# # scores2 = np.dot(cut_data, cut_data.T)[:, 0]
# #
# if True:
#     fig, ax = plt.subplots(3)
#     for i in xrange(no_measurements):
#         # print i, sum(u_full[i, :]), u_cut[i, :]
#         ax[0].plot(raw_data[i, :])
#         ax[1].plot(cut_data[i, :])
#         ax[2].plot(cut_data[i, :] - raw_data[i, :])
#     ax[0].set_title('full')
#     ax[1].set_title('PCA')
#     ax[2].set_title('difference')
#
# print cut_data[:, 0]
# if True:
#     if n == 2:
#         f7, (ax7, ax8, ax9) = plt.subplots(3)
#         ax7.plot(u_cut[:, 0], u_cut[:, 1], '-')
#         ax8.plot(v_cut[:, 0], v_cut[:, 1], '-')
#         ax9.plot(cut_data[:, 0]/sum(v_cut[:, 1]), cut_data[:, 1]/sum(v_cut[:, 0]), '-')
#         ax9.plot(cut_data[:, 0], cut_data[:, 1], '-')
#         ax7.set_xlabel('C1'); ax7.set_ylabel('C2')
#         ax8.set_xlabel('C1'); ax8.set_ylabel('C2')
#         ax9.set_xlabel('C1'); ax9.set_ylabel('C2')
#     elif n == 3:
#         from mpl_toolkits.mplot3d import Axes3D
#         f7 = plt.figure()
#         ax7 = f7.add_subplot(311, projection='3d')
#         ax8 = f7.add_subplot(312, projection='3d')
#         ax9 = f7.add_subplot(313, projection='3d')
#         ax7.plot(u_cut[:, 0], u_cut[:, 1], u_cut[:, 2])
#         ax8.plot(v_cut[:, 0], v_cut[:, 1], v_cut[:, 2])
#         # ax9.scatter(cut_data[:, 0]/sum(v_cut[:, 1]), cut_data[:, 1]/sum(v_cut[:, 0]), '-o')
#         ax9.plot(cut_data[:, 0], cut_data[:, 1], cut_data[:, 2])
#         ax7.set_xlabel('C1'); ax7.set_ylabel('C2'); ax7.set_zlabel('C3')
#         ax8.set_xlabel('C1'); ax8.set_ylabel('C2'); ax8.set_zlabel('C3')
#         ax9.set_xlabel('C1'); ax9.set_ylabel('C2'); ax9.set_zlabel('C3')
#
# f5, ax5 = plt.subplots(1)
# ax5.semilogx(tau, ideal_curve/max(ideal_curve), 'g--', label='model')
# ax5.semilogx(tau, spec_sum/max(spec_sum), 'm-', label='spectrum: Re')
# ax5.semilogx(tau, scores2/max(abs(scores2)), 'o', label='PCA')
# ax5.semilogx(tau, (scores-scores[-1])/(scores[0]-scores[-1]), 'yo', label='PCA_Oleg')
# plt.legend()
#
# plt.show()
