#!/usr/bin/env python
from __future__ import print_function
import numpy as np
import pytraj as pt

# local
from utils import fn


def test_radgyr_tensor():
    traj = pt.iterload(fn("tz2.truncoct.nc"), fn("tz2.truncoct.parm7"))

    d = pt.radgyr_tensor(traj, mask=':1-13', dtype='dict')
    expectect_tensor = [[
        18.26403420734941, 12.330493458000614, 15.326805252790127,
        -2.125659395015517, -8.632996672082086, 4.7252454067158975
    ], [
        19.246862113820757, 11.432920183634895, 15.489687898615074,
        -3.0267399828080603, -9.129524039907787, 4.5109070835284255
    ], [
        19.75385387541704, 10.836915486429934, 15.803415368373425,
        -3.441456117622487, -9.28446388556888, 4.1670219147678464
    ], [
        19.269059015181764, 13.042016327597267, 15.817716819622737,
        -3.710072422051309, -9.489039068401922, 5.2288271979672745
    ], [
        19.813786389085553, 11.80314746736244, 15.44020623457249,
        -3.502106891601859, -9.099872095876812, 4.6923151789825015
    ], [
        19.60077776416306, 12.609058958163347, 15.896052062712025,
        -3.6477265975369044, -9.698869709134073, 4.850778745624434
    ], [
        18.174738583179614, 12.74891714333484, 15.832179262508554,
        -3.5327010118759263, -8.664441526832949, 5.420192653400531
    ], [
        19.507715402036883, 11.7096354909941, 14.459002582255001,
        -3.3556288916094963, -8.441118169033068, 4.699660712555152
    ], [
        18.534885181089813, 11.609977647369599, 15.403693965901535,
        -3.0397336400024666, -8.90681248448317, 5.083930132028563
    ], [
        19.232621566893187, 12.027900886619847, 15.150377798541307,
        -3.525614491193894, -8.611207782564183, 5.491302107043754
    ]]

    key = 'RoG_00000[Tensor]'
    np.testing.assert_almost_equal(d.get(key), expectect_tensor)

    d2 = pt.radgyr_tensor(traj, mask=':1-13')
    np.testing.assert_almost_equal(d2[1], expectect_tensor)
    np.testing.assert_almost_equal(d2[0], pt.radgyr(traj, ':1-13'))


def test_radgyr_default():
    traj = pt.iterload(fn("tz2.truncoct.nc"), fn("tz2.truncoct.parm7"))
    state = pt.load_cpptraj_state("""
    radgyr :1-13
    """, traj)
    state.run()
    np.testing.assert_almost_equal(
        pt.radgyr(traj, ':1-13'), state.data['RoG_00001'].values)
