# Path: veidt/potential/processing.py
# def pool_from(structures, energies=None, forces=None, stresses=None):
#     """
#     Method to convert structures and their properties in to
#     datapool format.
#     Args:
#         structures ([Structure]): The list of Pymatgen Structure object.
#         energies ([float]): The list of total energies of each structure
#             in structures list.
#         forces ([np.array]): List of (m, 3) forces array of each structure
#             with m atoms in structures list. m can be varied with each
#             single structure case.
#         stresses (list): List of (6, ) virial stresses of each
#             structure in structures list.
#     Returns:
#         ([dict])
#     """
#     energies = energies if energies else [None] * len(structures)
#     forces = forces if forces else [None] * len(structures)
#     stresses = stresses if stresses else [None] * len(structures)
#     datapool = [doc_from(structure, energy, force, stress)
#                 for structure, energy, force, stress
#                 in zip(structures, energies, forces, stresses)]
#     return datapool
# 
# def convert_docs(docs, include_stress=False, **kwargs):
#     """
#     Method to convert a list of docs into objects, e.g.,
#     Structure and DataFrame.
#     Args:
#         docs ([dict]): List of docs. Each doc should have the same
#             format as one returned from .dft.parse_dir.
#         include_stress (bool): Whether to include stress.
#     Returns:
#         A list of structures, and a DataFrame with energy and force
#         data in 'y_orig' column, data type ('energy' or 'force') in
#         'dtype' column, No. of atoms in 'n' column sharing the same row
#         of energy data while 'n' being 1 for the rows of force data.
#     """
#     structures, y_orig, n, dtype = [], [], [], []
#     for d in docs:
#         if isinstance(d['structure'], dict):
#             structure = Structure.from_dict(d['structure'])
#         else:
#             structure = d['structure']
#         outputs = d['outputs']
#         force_arr = np.array(outputs['forces'])
#         assert force_arr.shape == (len(structure), 3), \
#             'Wrong force array not matching structure'
#         force_arr = force_arr.ravel()
# 
#         if include_stress:
#             stress_arr = np.array(outputs['virial_stress'])
#             y = np.concatenate(([outputs['energy']], force_arr, stress_arr))
#             n.append(np.insert(np.ones(len(y) - 1), 0, d['num_atoms']))
#             dtype.extend(['energy'] + ['force'] * len(force_arr) + ['stress'] * 6)
#         else:
#             y = np.concatenate(([outputs['energy']], force_arr))
#             n.append(np.insert(np.ones(len(y) - 1), 0, d['num_atoms']))
#             dtype.extend(['energy'] + ['force'] * len(force_arr))
#         y_orig.append(y)
#         structures.append(structure)
#     df = pd.DataFrame(dict(y_orig=np.concatenate(y_orig), n=np.concatenate(n),
#                            dtype=dtype))
#     for k, v in kwargs.items():
#         df[k] = v
#     return structures, df


import os
import shutil
import tempfile
import unittest
import numpy as np
from monty.serialization import loadfn
from veidt.potential.processing import pool_from, convert_docs

CWD = os.getcwd()
test_datapool = loadfn(os.path.join(os.path.dirname(__file__), 'datapool.json'))

class PorcessingTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.this_dir = os.path.dirname(os.path.abspath(__file__))
        cls.test_dir = tempfile.mkdtemp()
        os.chdir(cls.test_dir)

    @classmethod
    def tearDownClass(cls):
        os.chdir(CWD)
        shutil.rmtree(cls.test_dir)

    def setUp(self):
        self.test_pool = test_datapool
        self.test_structures = []
        self.test_energies = []
        self.test_forces = []
        self.test_stresses = []
        for d in self.test_pool:
            self.test_structures.append(d['structure'])
            self.test_energies.append(d['outputs']['energy'])
            self.test_forces.append(d['outputs']['forces'])
            self.test_stresses.append(d['outputs']['virial_stress'])

    def test_pool_from(self):

# Path: veidt/potential/processing.py
# def pool_from(structures, energies=None, forces=None, stresses=None):
#     """
#     Method to convert structures and their properties in to
#     datapool format.
#     Args:
#         structures ([Structure]): The list of Pymatgen Structure object.
#         energies ([float]): The list of total energies of each structure
#             in structures list.
#         forces ([np.array]): List of (m, 3) forces array of each structure
#             with m atoms in structures list. m can be varied with each
#             single structure case.
#         stresses (list): List of (6, ) virial stresses of each
#             structure in structures list.
#     Returns:
#         ([dict])
#     """
#     energies = energies if energies else [None] * len(structures)
#     forces = forces if forces else [None] * len(structures)
#     stresses = stresses if stresses else [None] * len(structures)
#     datapool = [doc_from(structure, energy, force, stress)
#                 for structure, energy, force, stress
#                 in zip(structures, energies, forces, stresses)]
#     return datapool
# 
# def convert_docs(docs, include_stress=False, **kwargs):
#     """
#     Method to convert a list of docs into objects, e.g.,
#     Structure and DataFrame.
#     Args:
#         docs ([dict]): List of docs. Each doc should have the same
#             format as one returned from .dft.parse_dir.
#         include_stress (bool): Whether to include stress.
#     Returns:
#         A list of structures, and a DataFrame with energy and force
#         data in 'y_orig' column, data type ('energy' or 'force') in
#         'dtype' column, No. of atoms in 'n' column sharing the same row
#         of energy data while 'n' being 1 for the rows of force data.
#     """
#     structures, y_orig, n, dtype = [], [], [], []
#     for d in docs:
#         if isinstance(d['structure'], dict):
#             structure = Structure.from_dict(d['structure'])
#         else:
#             structure = d['structure']
#         outputs = d['outputs']
#         force_arr = np.array(outputs['forces'])
#         assert force_arr.shape == (len(structure), 3), \
#             'Wrong force array not matching structure'
#         force_arr = force_arr.ravel()
# 
#         if include_stress:
#             stress_arr = np.array(outputs['virial_stress'])
#             y = np.concatenate(([outputs['energy']], force_arr, stress_arr))
#             n.append(np.insert(np.ones(len(y) - 1), 0, d['num_atoms']))
#             dtype.extend(['energy'] + ['force'] * len(force_arr) + ['stress'] * 6)
#         else:
#             y = np.concatenate(([outputs['energy']], force_arr))
#             n.append(np.insert(np.ones(len(y) - 1), 0, d['num_atoms']))
#             dtype.extend(['energy'] + ['force'] * len(force_arr))
#         y_orig.append(y)
#         structures.append(structure)
#     df = pd.DataFrame(dict(y_orig=np.concatenate(y_orig), n=np.concatenate(n),
#                            dtype=dtype))
#     for k, v in kwargs.items():
#         df[k] = v
#     return structures, df


# Path: veidt/potential/tests/test_processing.py
import os
import shutil
import tempfile
import unittest
import numpy as np
from monty.serialization import loadfn
from veidt.potential.processing import pool_from, convert_docs


CWD = os.getcwd()
test_datapool = loadfn(os.path.join(os.path.dirname(__file__), 'datapool.json'))

class PorcessingTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.this_dir = os.path.dirname(os.path.abspath(__file__))
        cls.test_dir = tempfile.mkdtemp()
        os.chdir(cls.test_dir)

    @classmethod
    def tearDownClass(cls):
        os.chdir(CWD)
        shutil.rmtree(cls.test_dir)

    def setUp(self):
        self.test_pool = test_datapool
        self.test_structures = []
        self.test_energies = []
        self.test_forces = []
        self.test_stresses = []
        for d in self.test_pool:
            self.test_structures.append(d['structure'])
            self.test_energies.append(d['outputs']['energy'])
            self.test_forces.append(d['outputs']['forces'])
            self.test_stresses.append(d['outputs']['virial_stress'])

    def test_pool_from(self):

        test_pool = pool_from(self.test_structures, self.test_energies,