import scipy.io as sio
from ..tools.tools import nextpow2, flat_eq, here


def cum3x(x, y, z, maxlag=0, nsamp=0, overlap=0, flag='biased', k1=0):
    """
    Estimate the third-order cross-cumulant for a fixed lag.

    Args:
        x: data vector or matrix
        y: data vector or matrix (same dimensions as x)
        z: data vector or matrix (same dimensions as x)
            if x, y, and z are matrices, columns correspond to independent
            realizations; overlap is set to 0 and nsamp is set to the
            row dimension
        maxlag: maximum lag to be computed [default = 0]
        nsamp: samples per segment [default = data length]
        overlap: percentage overlap of segments [default = 0]
            overlap is clipped to the allowed range of [0, 99]
        flag: cumulant estimation flag
            'biased': biased estimates are computed [default]
            'unbiased': unbiased estimates are computed
        k1: fixed lag in the third-order cross-cumulant C3(m, k1) [default = 0]

    Returns:
        y_cum: estimated third-order cross-cumulant sequence
            E[x*(n) y(n + m) z(n + k1)], where -maxlag <= m <= maxlag
    """

    (lx, nrecs) = x.shape
    if (lx, nrecs) != y.shape or (lx, nrecs) != z.shape:
        raise ValueError('x,y,z should have identical dimensions')

    if lx == 1:
        lx = nrecs
        nrecs = 1

    if maxlag < 0: raise ValueError('"maxlag" must be non-negative')
    if nrecs > 1: nsamp = lx
    if nsamp <= 0 or nsamp > lx: nsamp = lx
    if nrecs > 1: overlap = 0
    overlap = max(0, min(overlap, 99))

    overlap = np.fix(overlap / 100 * nsamp)
    nadvance = nsamp - overlap

    if nrecs == 1:
        nrecs = np.fix((lx - overlap) / nadvance)

    nlags = 2 * maxlag + 1
    zlag = maxlag
    y_cum = np.zeros([nlags, 1])

    if flag == 'biased':
        scale = np.ones([nlags, 1]) / nsamp
    else:
        lsamp = lx - abs(k1)
        scale = make_arr((range(lsamp - maxlag, lsamp + 1), range(lsamp - 1, lsamp - maxlag - 1, -1)), axis=1).T
        scale = np.ones([2 * maxlag + 1, 1]) / scale

    if k1 >= 0:
        indx = np.arange(nsamp - k1).T
        indz = np.arange(k1, nsamp).T
    else:
        indx = np.arange(-k1, nsamp).T
        indz = np.arange(nsamp + k1)

    ind = np.arange(nsamp)

    for k in np.arange(nrecs):
        xs = x[ind]
        xs = xs - np.mean(xs)
        ys = y[ind]
        ys = ys - np.mean(ys)
        zs = z[ind]
        zs = zs - np.mean(zs)
        zs = np.conj(zs)

        u = np.zeros([nsamp, 1])
        u[indx] = xs[indx] * zs[indz]

        y_cum[zlag] = y_cum[zlag] + np.dot(u.T, ys)

        for m in np.arange(1, maxlag + 1):
            y_cum[zlag - m] = y_cum[zlag - m] + np.dot(u[m:nsamp].T, ys[0:nsamp - m])
            y_cum[zlag + m] = y_cum[zlag + m] + np.dot(u[0:nsamp - m].T, ys[m:nsamp])

        ind = ind + int(nadvance)

    y_cum = y_cum * scale / nrecs

    return y_cum


def test():
    y = sio.loadmat(here(__file__) + '/demo/ma1.mat')['y']

    # The right results are:
    #           "biased": [0.36338   0.42762   0.77703   0.84322   0.73021  -0.13123  -0.40743]
    #           "unbiased": [0.035591   0.041841   0.075956   0.082345   0.071379  -0.012840  -0.039905]
    print(cum3x(y, y, y, 3, 100, 0, "biased"))
    print(cum3x(y, y, y, 3, 100, 0, "unbiased"))


if __name__ == '__main__':
    test()
