import numpy as np
import pandas as pd
import random
import matplotlib.pyplot as plt


def gen_random_wave(all_times, A, T):
    signal = A * np.sin(2*np.pi*all_times / T)
    signal += np.random.normal(0, A/20, size=all_times.size)
    return pd.DataFrame({'time': all_times, 'signal': signal})

def sample_wave(df, samples):
    ids = random.sample(range(len(df)), samples)
    return df.loc[ids, :]

def habitual_sample_wave(df, sample_period, num_samples):
    SHIFT = 0.3
    times = df.time
    prop = np.sin(2 * np.pi / sample_period * times) + (1 - 2 * SHIFT)
    prop = np.clip(prop, 0, 1)
    ftimes = times[prop > 0.5]
    chosen_times = np.random.choice(ftimes, num_samples, replace=False)
    return df[df.time.isin(chosen_times)]


ts = np.arange(0, 100, 0.001)
w1 = gen_random_wave(ts, 5, 8)
w1s = habitual_sample_wave(w1, 2.3, 1000)
# w1s = sample_wave(w1, 500)

w1s.to_csv('lombscargle_demo.csv', index=False)


w2 = gen_random_wave(ts, 10, 10)
w3 = gen_random_wave(ts, 5, 7)
w4 = gen_random_wave(ts, 3, 3)
w5 = pd.DataFrame({'time': w2.time, 'signal': w2.signal + w3.signal + w4.signal})
w5s = habitual_sample_wave(w5, 2.3, 1000)

w5s.to_csv('lombscargle_activity.csv', index=False)



plt.plot(w1s.time, w1s.signal, '.')
plt.show()