from __future__ import division, absolute_import, print_function

import pytest
from numpy.testing import assert_almost_equal, assert_equal

import numpy as np

import gromacs.fileformats.xpm
from gromacs.formats import XPM

from ..datafiles import datafile

@pytest.fixture
def ssfile():
    return datafile("fileformats/ss.xpm.bz2")

@pytest.fixture
def xpm(ssfile):
    return XPM(filename=ssfile)

class TestXPM(object):
    def _run_tests(self, x):
        assert_equal(x.array.shape, (500, 769))
        assert_equal(x.xvalues, np.arange(0, 500))
        assert_equal(x.yvalues, np.arange(1, 770))
        assert_equal((x.array == 'A-Helix').sum(), 292829)

    def test_constructor(self, xpm):
        self._run_tests(xpm)

    def test_read(self, ssfile):
        x = XPM()
        x.read(ssfile)
        self._run_tests(x)

    def test_reversed_by_default(self, xpm):
        assert xpm.reverse

    def test_to_pd(self, xpm):
        df = xpm.to_df()
        assert_equal(df.shape, (500, 770))

