from twBase import *  # NOQA
import numpy as np  # NOQA
import pandas as pd  # NOQA

from pprint import pprint  # NOQA

T = 10
n = 10


def create_data(T=10, n=10, split=0.8):
    x = np.zeros((n, 2, T))
    idx_arr = np.arange(T)

    for i in range(n):
        idx = np.random.choice(idx_arr, 2, replace=False)
        x[i, 0, idx] = 1
        x[i, 1, :] = np.random.uniform(0, 1, size=T)

    #y = x[:, 1, :][x[:, 0, :] > 0].reshape((n, count))
    #y = y.sum(axis=1)

    y = (x[:,0,:]*x[:,1,:]).sum(axis=1)  # NOQA
    row = round(n*split)

    return x[:row], y[:row], x[row:], y[row:]

if __name__ == '__main__':
    pprint(create_data())

    sys.exit()
