import numpy as np
from sys import stdout

"""
Generating mock data to test out function approximation methods.

Parameter space is 9-dimensional.

M: peaks around 5e5
q: linear then plateau
a: linear-ish increase (decrease at high M)
e: linear decrease
Y0: linear increase

qS, qK: (sin(x)+1.25)**2
phiS: (sin(2x)+1.25)**2

tpl: log increase

We don't really care what the normalisation of this distribution is but it should be 10 +- 1 order of mag 
across most of it. We can fix it such that the M peak has an SNR of ~200 and that should sort it out.


"""


def fM(lM):
    return 10 * np.exp(-2 * (lM - 6) ** 2) + 0.1


def fq(lq):
    #lq = np.log10(q)
    return 3 / (lq - 4.1)


def fa(a):
    return 0.2*a + 0.9


def fe(e):
    return 1.1 - 0.2 * e


def fY(Y):
    return Y


def fTh(th, phase):
    return 0.4*(np.sin(2 * np.pi * th + phase)) ** 2 + 0.6


def fPh(ph, phase):
    return 0.4*(np.sin(2 * np.pi * 2 * ph + phase) + 1) ** 2 + 0.6


def ft(t):
    return t**0.5/4**0.5


def func_out(M, q, a, e, Y, qS, phiS, qK, tpl):
    # return fM(M) + fq(q) + fa(a) + fe(e) + fY(Y) + fTh(qS, np.pi/3) + fTh(qK, np.pi/4) + fPh(phiS, 3*np.pi/8) + ft(tpl)
    return fM(M) * fq(q) * fa(a) * fe(e) * fY(Y) * fTh(qS, np.pi / 3) * fTh(qK, np.pi / 4) * fPh(phiS,
                                                                                                 3 * np.pi / 8) * ft(
        tpl)


def grid_sampling(grid_dims, ranges):

    total = np.prod(grid_dims)
    print(f'Total number of points: {total:.3e}')


def cartesian_product_recursive(arrays, out=None):
    arrays = [np.asarray(x) for x in arrays]
    dtype = arrays[0].dtype

    n = np.prod([x.size for x in arrays])
    if out is None:
        out = np.zeros([n, len(arrays)], dtype=dtype)

    m = n // arrays[0].size
    out[:, 0] = np.repeat(arrays[0], m)
    if arrays[1:]:
        cartesian_product_recursive(arrays[1:], out=out[0:m, 1:])
        for j in range(1, arrays[0].size):
            out[j * m: (j + 1) * m, 1:] = out[0:m, 1:]
    return out


def get_data(grid_dims, sample_ranges, sample_size=None, test_size=None):
    # get grid data, sampled data and some test data

    print('Generating grid data')

    iterators = [np.linspace(nums[0], nums[1], grid_dims[i]) for i,nums in enumerate(sample_ranges)]
    total = np.prod(grid_dims)

    combinations = cartesian_product_recursive(iterators)

    xgrid = np.array(combinations)
    ygrid = np.zeros(xgrid.shape[0])

    for i, combo in enumerate(combinations):
        ygrid[i] = func_out(*combo.tolist())

    print('Generating sample data')
    unset_samp = 0
    if sample_size is None:
        unset_samp = 1
        sample_size = total
    xsamp = np.zeros((sample_size, len(grid_dims)))

    for i in range(len(grid_dims)):
        xsamp[:, i] = np.random.uniform(sample_ranges[i, 0], sample_ranges[i, 1], sample_size)

    ysamp = np.zeros(sample_size)

    for i, row in enumerate(xsamp):
        ysamp[i] = func_out(*row.tolist())

    print('Generating test data')

    if test_size is None:
        if unset_samp:
            test_size = total
        else:
            test_size = sample_size

    xtest = np.zeros((test_size, len(grid_dims)))

    for i in range(len(grid_dims)):
        xtest[:, i] = np.random.uniform(sample_ranges[i, 0], sample_ranges[i, 1], test_size)

    ytest = np.zeros(test_size)

    for i, row in enumerate(xtest):
        ytest[i] = func_out(*row.tolist())

    return xgrid, ygrid, xsamp, ysamp, xtest, ytest


if __name__ == '__main__':
    # nM = 12
    # nq = 8
    # na = 7
    # ne = 7
    # nY = 7
    # nTh = 7
    # nPh = 8
    # nt = 8

    nM = 8
    nq = 5
    na = 4
    ne = 4
    nY = 4
    nTh = 5
    nPh = 8
    nt = 5

    # Our prior space (params are M,q,a,e,Y,thetaS,phiS,thetaK,t)
    sample_ranges = np.array([
        [np.log10(8e4), np.log10(5e7)],
        [np.log10(5e4), np.log10(5e7)],
        [0.01, 0.99],
        [0.01, 0.5],
        [0.5, 0.99],
        [0, np.pi],
        [0, 2 * np.pi],
        [0, np.pi],
        [0.5, 4],
    ])

    xgrid, ygrid, xsamp, ysamp, xtest, ytest = get_data([nM,nq,na,ne,nY,nTh,nPh,nPh,nt],sample_ranges, sample_size=int(1e6))

    fp = '../_data/{}'
    np.save(fp.format('xgrid.npy'),xgrid)
    np.save(fp.format('ygrid.npy'),ygrid)
    np.save(fp.format('xsamp.npy'),xsamp)
    np.save(fp.format('ysamp.npy'),ysamp)
    np.save(fp.format('xtest.npy'),xtest)
    np.save(fp.format('ytest.npy'),ytest)

    print('Data generation complete.')