"""
make sine initial conditions and times
"""

import torch
import argparse
from numpy import pi

parser = argparse.ArgumentParser()
parser.add_argument('--train', type=int, choices=[0, 1], default=0)
parser.add_argument('--ntimes', type=int, choices=[10, 50], default=50)
args = parser.parse_args()




train = args.train
ntimes = args.ntimes

if train:
    nsines = 15
    name = 'train'
else:
    nsines = 5
    name = 'val'
    
    
xv = torch.empty((nsines, 2))
t = torch.empty((nsines, ntimes, 1))


for i in range(nsines):
    initial_cond = 2*torch.rand(2)-1
    xv[i] = initial_cond
    times, _ = (2*pi*torch.rand(ntimes)).sort()
    times[0] = 0
    times = times.unsqueeze(-1)
    t[i] = times
    
xv = xv.float()
t = t.float()

torch.save(xv, 'sines_xv0_'+str(ntimes)+'_'+name+'.pt')
torch.save(t, 'sines_t_'+str(ntimes)+'_'+name+'.pt')

