import helperFuncs as hp
import numpy as np
from molSimplify.python_nn.ANN import simple_network_builder


def test_old_ANN(tmpdir):
    testName = "old_ANN"
    threshMLBL = 0.1
    threshLG = 1.0
    threshOG = 2.0
    [passNumAtoms, passMLBL, passLG, passOG, pass_report, _] = hp.runtest(
        tmpdir, testName, threshMLBL, threshLG, threshOG)
    assert passNumAtoms
    assert passMLBL
    assert passLG
    assert passOG
    assert pass_report


def test_simple_network_builder():
    """This test was created by capturing the inputs and outputs generated by
    the old pybrain implementation during the test_old_ANN() test function."""
    split_input = np.array(
        [-0.4917839, -0.56726225, 1.61974697, -0.50082508, -0.30476799,
         -0.90272848, 0.49027265, 0.65977966, 0.65377293, -0.5926822,
         -0.86257212, -0.25131881, -0.69807754, 1.76535323, -0.18234301,
         -0.16467918, -0.82545825, 1.43132271, -0.12278255, 1.52889381,
         1.11584538, -0.96668963, -1.05525312, -0.95836809, -1.25775219])
    split_ref_output = np.array([-0.90325581])

    n = simple_network_builder([25, 50, 50], "ms_split")
    output = n.activate(split_input)
    np.testing.assert_allclose(output, split_ref_output)

    # Same for slope
    slope_input = np.array(
        [-0.48069367, -0.551563, 1.64555217, -0.52515453, -0.31631299,
         -0.91503729, 0.67783389, 0.66328568, -0.6130196, -0.87338511,
         -0.25824387, -0.66613123, -0.57789148, 4.88496234, -0.15250621,
         -0.81416039, -0.69302998, 7.55021286, 1.54277755, 1.11294416,
         -0.96408929, -1.09251887, -0.97338522, -1.26357339])
    slope_ref_output = np.array([0.64270068])
    n = simple_network_builder([24, 50, 50], "ms_slope")
    output = n.activate(slope_input)
    np.testing.assert_allclose(output, slope_ref_output)

    ls_input = np.array(
        [-0.48919416, -0.57361705, 1.62646982, -0.49824537, -0.30638216,
         -0.92095277, 0.51297477, 0.66726202, 0.66624771, -0.57870037,
         -0.84940138, -0.24190509, -0.69676993, 1.75745942, -0.1917933,
         -0.16730346, -0.81139387, 1.43670349, -0.13278737,
         1.53683663, 1.11609907, -1.03399044, -0.82734353,
         -0.96566302, -1.26085904])
    ls_ref_output = np.array([-0.65810288])
    n = simple_network_builder([25, 50, 50], "ms_ls")
    output = n.activate(ls_input)
    np.testing.assert_allclose(output, ls_ref_output)

    hs_input = np.array(
        [-0.49953556, -0.56090563, 1.650425, -0.51368646, -0.30152256,
         -0.8952212, 0.5020875, 0.65942027, 0.65335421, -0.56599111,
         -0.84425765, -0.23805921, -0.69938591, 1.81009507, -0.18946893,
         -0.15931486, -0.83967322, 1.48102596, -0.12278255, 1.53846941,
         1.11789832, -1.08058748, -0.87232256, -0.94605866, -1.24373718])
    hs_ref_output = np.array([-0.54344953])
    n = simple_network_builder([25, 50, 50], "ms_hs")
    output = n.activate(hs_input)
    np.testing.assert_allclose(output, hs_ref_output)
