import numpy as np

from tqdm import tqdm
from plipy.dpdpl import DeepPrimalDualPriorLearningConv1d
from pathlib import Path


RESULTS = Path(__file__).resolve().parents[1] / "results"
N_TRAINING = 20


def theta(t, theta0, w0):
    return theta0 * np.sin(w0 * t)


def learning_1d(freq, sampling_rate, time=10):

    dt = 1 / sampling_rate
    N = 1000
    theta0 = np.random.random(N)
    T = np.arange(0, time, dt)

    xx, yy = np.meshgrid(T, theta0)

    thetas = theta(xx, yy, 2 * np.pi * freq)
    noise = np.random.normal(scale=thetas.std() / 5, size=thetas.shape)
    thetas_noisy = thetas + noise

    results = []

    while len(results) != N_TRAINING:
        dpdpl_ed = DeepPrimalDualPriorLearningConv1d(n_components=1,
                                                     n_iter=10,
                                                     lambd=0.5,
                                                     kernel_size=3,
                                                     device="cuda:1")
        dpdpl_ed.fit(thetas_noisy.T)

        kernel_ed = dpdpl_ed.get_prior()[0][0]
        kernel_ed / kernel_ed[0]
        w = np.sqrt(2 - (- kernel_ed / kernel_ed[0])[1]) / dt

        if not np.isnan(w):
            results.append(w / (2 * np.pi))

    return np.array(results).mean()


freq = np.arange(1, 11, 1)
sampling_rate = np.arange(10, 110, 10)
results = np.zeros((len(freq), len(sampling_rate)))

for i in tqdm(range(len(freq))):
    for j in range(len(sampling_rate)):
        print(sampling_rate[j])
        w_result = learning_1d(freq[i], sampling_rate[j])
        results[i, j] = w_result

np.save(str(RESULTS / "pde_learning_freq.npy"), freq)
np.save(str(RESULTS / "pde_learning_sampling.npy"), sampling_rate)
np.save(str(RESULTS / "pde_learning_results.npy"), results)
