import numpy as np
from treelas import post_order, TreeInstance


def test_demo_3x7_postord():
    parent = np.array([0, 4, 5, 0, 3, 4, 7, 8, 5, 6, 7, 8,
                       9, 14, 17, 12, 15, 16, 19, 16, 17])
    po = post_order(parent, include_root=True)
    expect = np.array([12, 11, 19, 20, 21, 14, 15, 18, 17, 16, 13,
                       10, 7, 8, 9, 3, 6, 2, 5, 4, 1], dtype='i4') - 1
    assert (po == expect).all()


def test_demo_3x7():
    y = np.fromstring("0.62 0.73 0.71 1.5 1.17 0.43 1.08 0.62 " +
                      "1.73 0.95 1.46 1.6 1.16 0.38 0.9 0.32 " +
                      "-0.48 0.95 1.08 0.02 0.4", sep=" ")
    parent = np.array([0, 4, 5, 0, 3, 4, 7, 8, 5, 6, 7, 8,
                       9, 14, 17, 12, 15, 16, 19, 16, 17])
    lam = 1.0
    prob = TreeInstance(y, parent, lam=lam)
    assert prob.root == 0
    assert prob.parent.dtype == np.int32
    prob.solve()
    assert abs(prob.x.mean() - prob.y.mean()) < 1e-15
    assert len(np.unique(prob.x)) == 2
    assert max(np.abs(prob.dual[2:]) - lam) < 1e-12
    assert max(np.abs(prob.gamma)) < 1e-15

[((6, 13, 7, 58), 'numpy.array', 'np.array', ({(6, 22, 7, 57): '[0, 4, 5, 0, 3, 4, 7, 8, 5, 6, 7, 8, 9, 14, 17, 12, 15, 16, 19, 16, 17]'}, {}), '([0, 4, 5, 0, 3, 4, 7, 8, 5, 6, 7, 8, 9, 14, 17, 12, 15, 16, 19, 16,\n    17])', True, 'import numpy as np\n'), ((8, 9, 8, 46), 'treelas.post_order', 'post_order', (), '', False, 'from treelas import post_order, TreeInstance\n'), ((15, 8, 17, 58), 'numpy.fromstring', 'np.fromstring', (), '', True, 'import numpy as np\n'), ((18, 13, 19, 58), 'numpy.array', 'np.array', ({(18, 22, 19, 57): '[0, 4, 5, 0, 3, 4, 7, 8, 5, 6, 7, 8, 9, 14, 17, 12, 15, 16, 19, 16, 17]'}, {}), '([0, 4, 5, 0, 3, 4, 7, 8, 5, 6, 7, 8, 9, 14, 17, 12, 15, 16, 19, 16,\n    17])', True, 'import numpy as np\n'), ((21, 11, 21, 43), 'treelas.TreeInstance', 'TreeInstance', (), '', False, 'from treelas import post_order, TreeInstance\n'), ((9, 13, 10, 66), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((26, 15, 26, 32), 'numpy.unique', 'np.unique', ({(26, 25, 26, 31): 'prob.x'}, {}), '(prob.x)', True, 'import numpy as np\n'), ((28, 15, 28, 33), 'numpy.abs', 'np.abs', ({(28, 22, 28, 32): 'prob.gamma'}, {}), '(prob.gamma)', True, 'import numpy as np\n'), ((27, 15, 27, 36), 'numpy.abs', 'np.abs', ({(27, 22, 27, 35): 'prob.dual[2:]'}, {}), '(prob.dual[2:])', True, 'import numpy as np\n')]