import vaex
import numpy as np
from common import *

# def test_count_multiple_selections():

def test_sum(df, ds_trimmed):
    df.select("x < 5")
    np.testing.assert_array_almost_equal(df.sum("x", selection=None), np.nansum(ds_trimmed.data.x))
    np.testing.assert_array_almost_equal(df.sum("x", selection=True), np.nansum(ds_trimmed.data.x[:5]))
    np.testing.assert_array_almost_equal(df.sum(df.x, selection=None), np.nansum(ds_trimmed.data.x))
    np.testing.assert_array_almost_equal(df.sum(df.x, selection=True), np.nansum(ds_trimmed.data.x[:5]))

    df.select("x > 5")
    np.testing.assert_array_almost_equal(df.sum("m", selection=None), np.nansum(ds_trimmed.data.m))
    np.testing.assert_array_almost_equal(df.sum("m", selection=True), np.nansum(ds_trimmed.data.m[6:]))
    np.testing.assert_array_almost_equal(df.m.sum(selection=True), np.nansum(ds_trimmed.data.m[6:]))

    df.select("x < 5")
    # convert to float
    x = ds_trimmed.columns["x"]
    y = ds_trimmed.data.y
    x_with_nan = x * 1
    x_with_nan[0] = np.nan
    ds_trimmed.columns["x"] = x_with_nan
    np.testing.assert_array_almost_equal(df.sum("x", selection=None), np.nansum(x))
    np.testing.assert_array_almost_equal(df.sum("x", selection=True), np.nansum(x[:5]))

    task = df.sum("x", selection=True, delay=True)
    df.execute()
    np.testing.assert_array_almost_equal(task.get(), np.nansum(x[:5]))


    np.testing.assert_array_almost_equal(df.sum("x", selection=None, binby=["x"], limits=[0, 10], shape=1), [np.nansum(x)])
    np.testing.assert_array_almost_equal(df.sum("x", selection=True, binby=["x"], limits=[0, 10], shape=1), [np.nansum(x[:5])])
    np.testing.assert_array_almost_equal(df.sum("x", selection=None, binby=["y"], limits=[0, 9**2+1], shape=1), [np.nansum(x)])
    np.testing.assert_array_almost_equal(df.sum("x", selection=True, binby=["y"], limits=[0, 9**2+1], shape=1), [np.nansum(x[:5])])
    np.testing.assert_array_almost_equal(df.sum("x", selection=None, binby=["x"], limits=[0, 10], shape=2), [np.nansum(x[:5]), np.nansum(x[5:])])
    np.testing.assert_array_almost_equal(df.sum("x", selection=True, binby=["x"], limits=[0, 10], shape=2), [np.nansum(x[:5]), 0])

    i = 7
    np.testing.assert_array_almost_equal(df.sum("x", selection=None, binby=["y"], limits=[0, 9**2+1], shape=2), [np.nansum(x[:i]), np.nansum(x[i:])])
    np.testing.assert_array_almost_equal(df.sum("x", selection=True, binby=["y"], limits=[0, 9**2+1], shape=2), [np.nansum(x[:5]), 0])

    i = 5
    np.testing.assert_array_almost_equal(df.sum("y", selection=None, binby=["x"], limits=[0, 10], shape=2), [np.nansum(y[:i]), np.nansum(y[i:])])
    np.testing.assert_array_almost_equal(df.sum("y", selection=True, binby=["x"], limits=[0, 10], shape=2), [np.nansum(y[:5]), 0])


def test_correlation_basics(df_local):
    df = df_local  # TODO: why does this not work with remote?
    correlation = df.correlation(df.y, df.y)
    np.testing.assert_array_almost_equal(correlation, 1.0)

    correlation = df.correlation(df.y, -df.y)
    np.testing.assert_array_almost_equal(correlation, -1.0)


def test_correlation(df_local):
    df = df_local  # TODO: why does this not work with remote?

    # convert to float
    x = df.x.to_numpy() #self.dataset_local.columns["x"][self.zero_index:10] = self.dataset_local.columns["x"][:10] * 1.
    y = df.y.to_numpy()
    def correlation(x, y):
        c = np.cov([x, y], bias=1)
        return c[0,1] / (c[0,0] * c[1,1])**0.5

    np.testing.assert_array_almost_equal(df.correlation([["x", "y"], ["x", "x**2"]], selection=None), [correlation(x, y), correlation(x, x**2)])

    df.select("x < 5")
    np.testing.assert_array_almost_equal(df.correlation("x", "y", selection=None), correlation(x, y))
    np.testing.assert_array_almost_equal(df.correlation("x", "y", selection=True), correlation(x[:5], y[:5]))

    np.testing.assert_array_almost_equal(df.correlation("x", "y", selection=None), correlation(x, y))
    np.testing.assert_array_almost_equal(df.correlation("x", "y", selection=True), correlation(x[:5], y[:5]))

    task = df.correlation("x", "y", selection=True, delay=True)