
from __future__ import division

from nose.tools import (
    assert_almost_equal,
    assert_equal,
    assert_true,
    assert_raises
)

import numpy as np

from numpy import allclose

import dit
from dit.exceptions import ditException
from dit.math import LogOperations
ops = LogOperations(2)
log2 = ops.log
exp2 = ops.exp

from dit.math.aitchison import (alr, alr_inv, basis, closure, clr, clr_inv,
                                dist, ilr, ilr_inv, inner, norm, perturbation,
                                power, subcomposition, _gm, _log2_gm)

def assert_almost_equal_array(x, y):
    for i, j in zip(x.ravel(), y.ravel()):
        assert_almost_equal(i, j)

def assert_equal_shape(x, y):
    assert_equal(x.shape, y.shape)

def test_closure():
    # 1D
    x = np.array([1.0, 1.0])
    y = np.array([0.5, 0.5])
    cl_x = closure(x)
    assert_equal_shape(x, cl_x)
    assert_almost_equal_array(cl_x, y)

    # 2D with multiple rows
    x = np.array([[1.0, 1.0], [1.0, 2.0]])
    y = np.array([[0.5, 0.5], [1/3.0, 2/3.0]])
    cl_x = closure(x)
    assert_equal_shape(x, cl_x)
    assert_almost_equal_array(cl_x, y)

    # 2D with one row
    x = np.array([[1.0, 1.0]])
    y = np.array([[0.5, 0.5]])
    cl_x = closure(x)
    assert_equal_shape(x, cl_x)
    assert_almost_equal_array(cl_x, y)

def test_closure_invariance():
    x = np.array([1, 2, 3, 4]) / 10
    y = x*2
    z = x*.5
    xcl = closure(x)
    ycl = closure(y)
    zcl = closure(z)

    assert_almost_equal_array(xcl, ycl)
    assert_almost_equal_array(ycl, zcl)
    assert_almost_equal_array(zcl, xcl)

def test_closure_unnormalizable():
    """
    Test an unnormalizable vector.

    """
    x = np.zeros(3)
    assert_raises(ditException, closure, x)

def test_perturbation():
    # 1D
    x = np.array([1/3.0, 2/3.0])
    y = np.array([3/4.0, 1/4.0])
    pxy_ = np.array([3/5.0, 2/5.0])
    pxy = perturbation(x, y)
    assert_true(x.shape == pxy.shape)
    assert_true(allclose(pxy_, pxy))

    # 2D with multiple rows
    x = np.array([[1/3.0, 2/3.0], [3/4.0, 1/4.0]])
    y = np.array([[3/4.0, 1/4.0], [1/6.0, 5/6.0]])
    pxy_ = np.array([[3/5.0, 2/5.0], [3/8.0, 5/8.0]])
    pxy = perturbation(x, y)
    assert_true(x.shape == pxy.shape)
    assert_true(allclose(pxy_, pxy))

    # 2D with one row
    x = np.array([[1/3.0, 2/3.0]])
    y = np.array([[3/4.0, 1/4.0]])
    pxy_ = np.array([[3/5.0, 2/5.0]])
    pxy = perturbation(x, y)
    assert_true(x.shape == pxy.shape)
    assert_true(allclose(pxy_, pxy))

def test_power():
    # 1D
    x = np.array([1/3.0, 2/3.0])
    y = 3
    px_ = np.array([1/9.0, 8/9.0])
    px = power(x, y)
    assert_true(px_.shape == px.shape)
    assert_true(allclose(px_, px))
    x = np.array([1/4.0, 3/4.0])
    y = 3
    px_ = np.array([1/28.0, 27/28.0])
    px = power(x, y)
    assert_true(px_.shape == px.shape)
    assert_true(allclose(px_, px))

    # 2D with multiple rows
    x = np.array([[1/3.0, 2/3.0], [1/4.0, 3/4.0]])
    y = [3, 3]
    px_ = np.array([[1/9.0, 8/9.0], [1/28.0, 27/28.0]])
    px = power(x, y)
    assert_true(px_.shape == px.shape)
    assert_true(allclose(px_, px))

    # 2D with one row
    x = np.array([[1/3.0, 2/3.0]])
    y = [3]
    px_ = np.array([[1/9.0, 8/9.0]])
    px = power(x, y)
    assert_true(px_.shape == px.shape)
    assert_true(allclose(px_, px))

def test_inner():
    # 1D
    x = np.array([1/3.0, 2/3.0])
    y = np.array([1/4.0, 3/4.0])
    z = np.array([2/5.0, 3/5.0])

    xy_ = 0.79248125036057804
    xy = inner(x, y)
    yx = inner(y, x)
    assert_almost_equal(xy_, xy)
    assert_almost_equal(xy_, yx)

    yz_ = 0.46357181398555231
    yz = inner(y, z)
    zy = inner(z, y)
    assert_almost_equal(yz_, yz)
    assert_almost_equal(yz_, zy)

    xz_ = 0.29248125036057804
    xz = inner(x, z)
    zx = inner(z, x)
    assert_almost_equal(xz_, xz)
    assert_almost_equal(xz_, zx)

    # 2D with multiple rows
    xx = np.array([x, y, z])
    yy = np.array([y, z, x])
    xxyy = inner(xx, yy)
    xxyy_ = np.array([xy_, yz_, xz_])
    assert_true(allclose(xxyy_, xxyy))

    # 2D with multiple rows in x and single row in y
    yy_ = 1.2560530643461303
    xx = np.array([x, y, z])
    yy = np.array([y])
    xxyy = inner(xx, yy)
    xxyy_ = np.array([xy_, yy_, yz_])
    assert_true(allclose(xxyy_, xxyy))
    yyxx = inner(yy, xx)
    assert_true(allclose(xxyy_, yyxx))

    # 2D with single rows
    x2d = np.atleast_2d(x)
    y2d = np.atleast_2d(y)
    xy2d = inner(x2d, y2d)
    xy2d_ = np.array([xy_])
    assert_true(allclose(xy2d_, xy2d))

def test_clr():
    x = np.array([1/3.0, 2/3.0])
    y = np.array([1/4.0, 3/4.0])

    x_clr = [-0.5, 0.5]
    y_clr = [-0.79248125036057782, 0.79248125036057815]

    # 1D
    x1d_clr = clr(x)
    x1d_clr_ = np.array(x_clr)
    assert_equal(x1d_clr.shape, x1d_clr_.shape)
    assert_true(allclose(x1d_clr, x1d_clr_))

    # 2D
    xy = np.array([x, y])
    xy_clr = clr(xy)
    xy_clr_ = np.array([x_clr, y_clr])
    assert_equal(xy_clr.shape, xy_clr_.shape)
    assert_true(allclose(xy_clr, xy_clr_))

def test_alr():
    x = np.array([1/3.0, 2/3.0])
    y = np.array([1/4.0, 3/4.0])

    x_alr = [-1]
    y_alr = [-1.5849625007211563]

    # 1D
    x1d_alr = alr(x)
    x1d_alr_ = np.array(x_alr)