def test_order_1():
    with tempdir():
        with dataset('ds1.h5', 'w', 'hdf5') as ds1:
            for i in range(20):
                ds1[i] = np.random.randn(10)
            assert list(ds1.keys()) == list(range(20))

        with dataset('ds1/', 'w', 'dir-npy') as ds1:
            for i in range(20):
                ds1[i] = np.random.randn(10)
            assert list(ds1.keys()) == list(range(20))

def test_mdtraj_1():
    ds = dataset(get_fn('') + '*.pdb', fmt='mdtraj', verbose=True)
    print(ds.keys())
    print(ds.get(0))
    print(ds.provenance)

    ds = dataset(get_fn('') + '*.pdb', fmt='mdtraj', atom_indices=[1, 2],
                 verbose=True)
    print(ds.keys())
    print(ds.get(0))
    print(ds.provenance)

def test_union_3():
    with tempdir():
        # This doesn't work with py2.6
        with dataset('ds1/', 'w', 'dir-npy') as ds1, \
                dataset('ds2/', 'w', 'dir-npy') as ds2:
            ds1[0] = np.random.randn(10, 2)
            ds1[1] = np.random.randn(10)
            ds2[0] = np.random.randn(10,4)
            # Uneven length!

        with assert_raises(ValueError):
            mds = dataset(['ds1', 'ds2'])

def test_hdf5_3():
    with tempdir():
        with dataset('ds.h5', 'w', 'hdf5') as ds:
            ds[0] = np.random.randn(10)
            ds[1] = np.random.randn(10)
            ref_sum = _sum_helper(ds)

        iter_args = (dataset('ds.h5') for _ in range(5))

        sums = Parallel(n_jobs=2)(
            delayed(_sum_helper)(a) for a in iter_args)

        assert all(s == ref_sum for s in sums)

def test_uneven_n():
    with tempdir():
        # This doesn't work with py2.6
        with dataset('ds1/', 'w', 'dir-npy') as ds1, \
                dataset('ds2/', 'w', 'dir-npy') as ds2:
            ds1[0] = np.random.randn(10, 2)
            ds1[1] = np.random.randn(5, 2)
            ds2[0] = np.random.randn(10, 4)
            # Uneven number of trajs!

            fu = FeatureUnion(normalize=False)
            with assert_raises(ValueError):
                fu.fit((ds1, ds2))

def test_union_2():
    with tempdir():
        # This doesn't work with py2.6
        with dataset('ds1/', 'w', 'dir-npy') as ds1, \
                dataset('ds2/', 'w', 'dir-npy') as ds2:
            ds1[0] = np.random.randn(10, 2)
            ds1[1] = np.random.randn(10)
            ds2[0] = np.random.randn(10,4)
            ds2[1] = np.random.randn(10,4)


        mds = dataset(['ds1', 'ds2'])
        mds_out = mds.create_derived('derived', fmt='dir-npy')
        assert len(mds_out.provenance.split('\n')) > 0

    def test_assign(self):
        with open(os.devnull) as dn:
            subprocess.call(
                [
                    'msmb', 'SolventShellsAssigner', '--trjs', self.traj_fn,
                    '--solute_indices', self.ute_fn, '--solvent_indices',
                    self.vent_fn, '--n_shells', '3', '--shell_width', '1',
                    '--out', self.outfn, '--chunk', '2'
                ], stdout=dn, stderr=dn
            )

        data = dataset(self.outfn)[0]

        should_be = np.array([
            [0, 0, 0, 0],
            [0, 1, 0, 0],
            [1, 0, 0, 1],
            [1, 1, 0, 1],
            [2, 0, 0, 2],
            [2, 1, 0, 2],
            # 3
            # 4
            [5, 1, 0, 0],
            [6, 1, 0, 1],
            [7, 1, 0, 2],
            # 8
        ])

        np.testing.assert_array_equal(data, should_be)

    def test_partial_transform(self):
        with open(os.devnull) as dn:
            subprocess.call(
                [
                    'msmb', 'SolventShellsFeaturizer', '--trjs', self.traj_fn,
                    '--solute_indices', self.ute_fn, '--solvent_indices',
                    self.vent_fn, '--n_shells', '3', '--shell_width', '1',
                    '--out', self.outfn
                ], stdout=dn, stderr=dn
            )
        data = dataset(self.outfn)[0]

        norm = np.asarray([4 * np.pi * r ** 2 for r in [0.5, 1.5, 2.5]])
        should_be = np.array([
            [2, 0, 0],
            [0, 2, 0],
            [0, 0, 2],
            [0, 0, 0],
            [0, 0, 0],
            [1, 0, 0],
            [0, 1, 0],
            [0, 0, 1],
            [0, 0, 0]
        ]) / norm

        np.testing.assert_array_equal(data, should_be)

def test_1():
    path = tempfile.mkdtemp()
    shutil.rmtree(path)
    try:
        X = np.random.randn(10, 2)
        ds = dataset(path, 'w', 'dir-npy')
        ds[0] = X
        assert set(os.listdir(path)) == set(('PROVENANCE.txt', '00000000.npy'))
        np.testing.assert_array_equal(ds[0], X)

        assert_raises(IndexError, lambda: ds[1])
        assert len(ds) == 1

        Y = np.zeros((10, 1))
        Z = np.ones((2, 2))
        ds[1] = Y
        ds[2] = Z
        np.testing.assert_array_equal(ds[1], Y)
        np.testing.assert_array_equal(ds[2], Z)
        assert len(ds) == 3

        for i, item in enumerate(ds):
            np.testing.assert_array_equal(item, [X, Y, Z][i])
    except:
        raise
    finally:
        shutil.rmtree(path)

def test_2():
    path1 = tempfile.mkdtemp()
    path2 = tempfile.mkdtemp()
    shutil.rmtree(path1)
    shutil.rmtree(path2)
    try:

        X = np.random.randn(10, 2)
        Y = np.random.randn(10, 2)
        ds1 = dataset(path1, 'w', 'dir-npy')
        ds1[0] = X

        ds2 = ds1.create_derived(path2)
        ds2[0] = Y

        np.testing.assert_array_equal(ds1[0], X)
        np.testing.assert_array_equal(ds2[0], Y)
        assert len(ds1) == 1
        assert len(ds2) == 1

        prov2 = ds2.provenance
        print(prov2)
        assert 2 == sum([s.startswith('  Command') for s in prov2.splitlines()])

    except:
        raise
    finally:
        shutil.rmtree(path1)
        shutil.rmtree(path2)

    def start(self):
        if os.path.exists(self.out):
            self.error('File exists: %s' % self.out)

        print(self.instance)
        if os.path.exists(os.path.expanduser(self.top)):
            top = os.path.expanduser(self.top)
        else:
            top = None

        traj_dataset = MDTrajDataset(self.trjs, topology=top,
                                     stride=self.instance.stride, verbose=False)

        with dataset(self.assignments, mode='r') as assn_dataset:
            out_dataset = assn_dataset.create_derived(self.out, fmt='dir-npy')
            pbar = ProgressBar(widgets=[Percentage(), Bar(), ETA()],
                               maxval=len(assn_dataset)).start()
            for tr_key, as_key in pbar(
                    zip(traj_dataset.keys(), assn_dataset.keys())
            ):
                out_dataset[as_key] = self.instance.partial_transform(
                    (traj_dataset[tr_key], assn_dataset[as_key])
                )
            out_dataset.close()

        print("\nSaving transformed dataset to '%s'" % self.out)
        print("To load this dataset interactive inside an IPython")
        print("shell or notebook, run\n")
        print("  $ ipython")
        print("  >>> from msmbuilder.dataset import dataset")
        print("  >>> ds = dataset('%s')\n" % self.out)

def test_transform_command_1():
    with tempdir():
        shell("msmb KCenters -i {data_home}/alanine_dipeptide/*.dcd "
              "-o model.pkl --top {data_home}/alanine_dipeptide/ala2.pdb "
              "--metric rmsd".format(data_home=get_data_home()))
        shell("msmb TransformDataset -i {data_home}/alanine_dipeptide/*.dcd "
              "-m model.pkl -t transformed.h5 --top "
              "{data_home}/alanine_dipeptide/ala2.pdb".format(data_home=get_data_home()))

        eq(dataset('transformed.h5')[0], load('model.pkl').labels_[0])

def test_hdf5_1():
    with tempdir():
        ds = dataset('ds.h5', 'w', 'hdf5')
        print(ds.provenance)
        ds[0] = np.zeros(10)
        np.testing.assert_array_equal(ds.get(0), np.zeros(10))
        assert list(ds.keys()) == [0]
        assert len(ds) == 1

        ds[0] = np.random.randn(10, 1)
        ds[1] = np.random.randn(10, 2)
        ds[2] = np.random.randn(10, 3)

        np.testing.assert_array_equal(ds[:][0], ds[0])
        np.testing.assert_array_equal(ds[:][1], ds[1])
        np.testing.assert_array_equal(ds[:][2], ds[2])

        ds.close()
        with dataset('ds.h5') as ds:
            assert ds[0].shape == (10, 1)

def test_atom_pairs_featurizer():
    with tempdir():
        shell('msmb AtomIndices -o all.txt --all -d -p %s/alanine_dipeptide/ala2.pdb' % get_data_home()),
        shell("msmb AtomPairsFeaturizer --trjs '{data_home}/alanine_dipeptide/*.dcd'"
              " --out pairs --pair_indices all.txt"
              " --top {data_home}/alanine_dipeptide/ala2.pdb".format(
                  data_home=get_data_home()))
        ds = dataset('pairs')
        assert len(ds) == 10
        assert ds[0].shape[1] == len(np.loadtxt('all.txt')**2)
        print(ds.provenance)

def test_union():
    with tempdir():
        # This doesn't work with py2.6
        with dataset('ds1.h5', 'w', 'hdf5') as ds1, \
             dataset('ds2.h5', 'w', 'hdf5') as ds2:
            ds1[0] = np.random.randn(10, 2)
            ds1[1] = np.random.randn(10)
            ds2[0] = np.random.randn(10,4)
            ds2[1] = np.random.randn(10,4)

            # Compare row sums
            rs1 = np.sum(ds1[0], axis=1) + np.sum(ds2[0], axis=1)
            rs2 = ds1[1] + np.sum(ds2[1], axis=1)

        mds = dataset(['ds1.h5', 'ds2.h5'])

        assert len(mds) == 2
        assert mds[0].shape == (10, 6)
        assert mds[1].shape == (10, 5)
        np.testing.assert_array_almost_equal(np.sum(mds[0], axis=1), rs1)
        np.testing.assert_array_almost_equal(np.sum(mds[1], axis=1), rs2)

def test_superpose_featurizer():
    with tempdir():
        shell('msmb AtomIndices -o all.txt --all -a -p %s/alanine_dipeptide/ala2.pdb' % get_data_home()),
        shell("msmb SuperposeFeaturizer --trjs '{data_home}/alanine_dipeptide/*.dcd'"
              " --out distances --atom_indices all.txt"
              " --reference_traj {data_home}/alanine_dipeptide/ala2.pdb"
              " --top {data_home}/alanine_dipeptide/ala2.pdb".format(
                  data_home=get_data_home()))
        ds = dataset('distances')
        assert len(ds) == 10
        assert ds[0].shape[1] == len(np.loadtxt('all.txt'))
        print(ds.provenance)

def test_append_dirnpy():
    path = tempfile.mkdtemp()
    shutil.rmtree(path)
    try:
        with dataset(path, 'w', 'dir-npy') as ds:
            ds[0] = np.random.randn(10, 2)
        with dataset(path, 'a', 'dir-npy') as ds:
            ds[1] = np.random.randn(10, 2)
        with dataset(path, 'a', 'dir-npy') as ds:
            ds[2] = np.random.randn(10, 2)
        with dataset(path, 'a', 'dir-npy') as ds:
            # Overwrite
            ds[2] = np.random.randn(10, 2)

        np.testing.assert_array_equal(ds[:][0], ds[0])
        np.testing.assert_array_equal(ds[:][1], ds[1])
        np.testing.assert_array_equal(ds[:][2], ds[2])


    finally:
        shutil.rmtree(path)

def test_4():
    path = tempfile.mkdtemp()
    shutil.rmtree(path)
    try:
        ds = dataset(path, 'w', 'dir-npy')
        ds[0] = np.random.randn(10, 2)
        v = ds.get(0, mmap=True)
        assert isinstance(v, np.memmap)
        np.testing.assert_array_equal(ds[0], v)
        del v  # close the underlying file
    finally:
        shutil.rmtree(path)

def test_dataset():
    with tempdir():
        # This doesn't work with py2.6
        with dataset('ds1.h5', 'w', 'hdf5') as ds1, \
                dataset('ds2.h5', 'w', 'hdf5') as ds2:
            ds1[0] = np.random.randn(10, 2)
            ds1[1] = np.random.randn(5, 2)
            ds2[0] = np.random.randn(10, 4)
            ds2[1] = np.random.randn(5, 4)

            # Compare row sums
            rs1 = np.sum(ds1[0], axis=1) + np.sum(ds2[0], axis=1)
            rs2 = np.sum(ds1[1], axis=1) + np.sum(ds2[1], axis=1)

            fu = FeatureUnion(normalize=False)
            mds = fu.fit_transform((ds1, ds2))

        assert len(mds) == 2
        assert mds[0].shape == (10, 6)
        assert mds[1].shape == (5, 6)
        np.testing.assert_array_almost_equal(np.sum(mds[0], axis=1), rs1)
        np.testing.assert_array_almost_equal(np.sum(mds[1], axis=1), rs2)

def test_3():
    path = tempfile.mkdtemp()
    shutil.rmtree(path)
    try:
        ds = dataset(path, 'w', 'dir-npy')
        ds[0] = np.random.randn(10, 2)
        ds[1] = np.random.randn(10, 2)
        ds[2] = np.random.randn(10, 2)

        np.testing.assert_array_equal(ds[:][0], ds[0])
        np.testing.assert_array_equal(ds[:][1], ds[1])
        np.testing.assert_array_equal(ds[:][2], ds[2])


    finally:
        shutil.rmtree(path)