# Path: flare/env.py
# class AtomicEnvironment:
#     def __init__(self, structure: Structure, atom: int, cutoffs, cutoffs_mask=None):
#     def setup_mask(self, cutoffs_mask):
#     def compute_env(self):
#     def as_str(self) -> str:
#     def as_dict(self, include_structure: bool = False):
#     def from_dict(dictionary):
#     def from_file(
#         file_name: str, format: str = ""
#     ) -> Union["flare.env.AtomicEnvironment", List["flare.env.AtomicEnvironment"]]:
#     def __str__(self):
#     def force(self) -> "np.ndarray":
# 
# Path: flare/struc.py
# class Structure:
# class StructureSource(object):
# class ForceSource(object):
# class Trajectory(StructureSource, ForceSource):
#     def __init__(
#         self,
#         cell: "ndarray",
#         species: Union[List[str], List[int]],
#         positions: "ndarray",
#         mass_dict: dict = None,
#         prev_positions: "ndarray" = None,
#         species_labels: List[str] = None,
#         forces=None,
#         stds=None,
#         energy: float = None,
#     ):
#     def positions(self):
#     def wrapped_positions(self):
#     def positions(self, position_array):
#     def cell(self):
#     def vec1(self):
#     def vec2(self):
#     def vec3(self):
#     def cell_transpose(self):
#     def cell_transpose_inverse(self):
#     def cell_dot(self):
#     def cell_dot_inverse(self):
#     def cell(self, cell_array):
#     def get_cell_dot(cell_array):
#     def raw_to_relative(
#         positions: "ndarray", cell_transpose: "ndarray", cell_dot_inverse: "ndarray"
#     ) -> "ndarray":
#     def relative_to_raw(
#         relative_positions: "ndarray",
#         cell_transpose_inverse: "ndarray",
#         cell_dot: "ndarray",
#     ) -> "ndarray":
#     def wrap_positions(self) -> "ndarray":
#     def indices_of_specie(self, specie: Union[int, str]) -> List[int]:
#     def __str__(self) -> str:
#     def __len__(self) -> int:
#     def as_dict(self) -> dict:
#     def as_str(self) -> str:
#     def from_dict(dictionary: dict) -> "flare.struc.Structure":
#     def from_ase_atoms(atoms: "ase.Atoms", cell=None) -> "flare.struc.Structure":
#     def to_ase_atoms(self) -> "ase.Atoms":
#     def to_pmg_structure(self):
#     def from_pmg_structure(structure: "pymatgen Structure") -> "flare Structure":
#     def is_valid(self, tolerance: float = 0.5) -> bool:
#     def to_xyz(
#         self,
#         extended_xyz: bool = True,
#         print_stds: bool = False,
#         print_forces: bool = False,
#         print_max_stds: bool = False,
#         print_energies: bool = False,
#         predict_energy=None,
#         dft_forces=None,
#         dft_energy=None,
#         target_atoms=None,
#         timestep=-1,
#         write_file: str = "",
#         append: bool = False,
#     ) -> str:
#     def from_file(
#         file_name: str, format: str = "", as_trajectory: bool = False
#     ) -> Union["flare.struc.Structure", List["flare.struc.Structure"]]:
#     def __eq__(self, other):
#     def __init__(self):
#     def get_next_structure(self) -> Structure:
#     def write_file(self):
#     def __init__(self):
#     def get_next_force(self, *args, **kwargs) -> "np.ndarray":
#     def pre_force(self, *args, **kwargs):
#     def __init__(
#         self, frames: List[Structure] = None, iterate_strategy: Union[int, str] = 1
#     ):
#     def get_next_structure(self) -> Union[Structure, None]:
#     def cur_frame(self):
#     def cur_forces(self):
#     def get_next_force(self, index: int = -1) -> "np.ndarray":
#     def __getitem__(self, item):
#     def __len__(self):
#     def __next__(self):
#     def __iter__(self):
#     def append(self, frame: Structure):
# def get_unique_species(species: List[Any]) -> (List, List[int]):


import os, shutil, re
import numpy as np
from flare import env, struc
        gp_force = np.sum(gp_f, axis=0)
        gp_energy = np.sum(gp_e, axis=0)
        print("gp_e", gp_e)
        print("gp_f")
        print(gp_f)

        map_ind = mgp_model.find_map_index(spc)
        xyzs = np.zeros_like(xyzs)
        xyzs[:, 0] = np.ones_like(xyzs[:, 0])
        f, _, _, e = mgp_model.maps[map_ind].predict(lengths, xyzs)

        assert np.allclose(gp_force, f, rtol=1e-2)
        assert np.allclose(gp_energy, e, rtol=1e-2)


def get_triplet_env(r1, r2, r12, grid_env):
    grid_env.bond_array_3 = np.array([[r1, 1, 0, 0], [r2, 0, 0, 0]])
    grid_env.cross_bond_dists = np.array([[0, r12], [r12, 0]])
    print(grid_env.ctype, grid_env.etypes)

    return grid_env


def get_grid_env(GP, species, bodies):
    if isinstance(GP.cutoffs, dict):
        max_cut = np.max(list(GP.cutoffs.values()))
    else:
        max_cut = np.max(GP.cutoffs)
    big_cell = np.eye(3) * 100
    positions = [[(i + 1) / (bodies + 1) * 0.1, 0, 0] for i in range(bodies)]

# Path: flare/env.py
# class AtomicEnvironment:
#     def __init__(self, structure: Structure, atom: int, cutoffs, cutoffs_mask=None):
#     def setup_mask(self, cutoffs_mask):
#     def compute_env(self):
#     def as_str(self) -> str:
#     def as_dict(self, include_structure: bool = False):
#     def from_dict(dictionary):
#     def from_file(
#         file_name: str, format: str = ""
#     ) -> Union["flare.env.AtomicEnvironment", List["flare.env.AtomicEnvironment"]]:
#     def __str__(self):
#     def force(self) -> "np.ndarray":
# 
# Path: flare/struc.py
# class Structure:
# class StructureSource(object):
# class ForceSource(object):
# class Trajectory(StructureSource, ForceSource):
#     def __init__(
#         self,
#         cell: "ndarray",
#         species: Union[List[str], List[int]],
#         positions: "ndarray",
#         mass_dict: dict = None,
#         prev_positions: "ndarray" = None,
#         species_labels: List[str] = None,
#         forces=None,
#         stds=None,
#         energy: float = None,
#     ):
#     def positions(self):
#     def wrapped_positions(self):
#     def positions(self, position_array):
#     def cell(self):
#     def vec1(self):
#     def vec2(self):
#     def vec3(self):
#     def cell_transpose(self):
#     def cell_transpose_inverse(self):
#     def cell_dot(self):
#     def cell_dot_inverse(self):
#     def cell(self, cell_array):
#     def get_cell_dot(cell_array):
#     def raw_to_relative(
#         positions: "ndarray", cell_transpose: "ndarray", cell_dot_inverse: "ndarray"
#     ) -> "ndarray":
#     def relative_to_raw(
#         relative_positions: "ndarray",
#         cell_transpose_inverse: "ndarray",
#         cell_dot: "ndarray",
#     ) -> "ndarray":
#     def wrap_positions(self) -> "ndarray":
#     def indices_of_specie(self, specie: Union[int, str]) -> List[int]:
#     def __str__(self) -> str:
#     def __len__(self) -> int:
#     def as_dict(self) -> dict:
#     def as_str(self) -> str:
#     def from_dict(dictionary: dict) -> "flare.struc.Structure":
#     def from_ase_atoms(atoms: "ase.Atoms", cell=None) -> "flare.struc.Structure":
#     def to_ase_atoms(self) -> "ase.Atoms":
#     def to_pmg_structure(self):
#     def from_pmg_structure(structure: "pymatgen Structure") -> "flare Structure":
#     def is_valid(self, tolerance: float = 0.5) -> bool:
#     def to_xyz(
#         self,
#         extended_xyz: bool = True,
#         print_stds: bool = False,
#         print_forces: bool = False,
#         print_max_stds: bool = False,
#         print_energies: bool = False,
#         predict_energy=None,
#         dft_forces=None,
#         dft_energy=None,
#         target_atoms=None,
#         timestep=-1,
#         write_file: str = "",
#         append: bool = False,
#     ) -> str:
#     def from_file(
#         file_name: str, format: str = "", as_trajectory: bool = False
#     ) -> Union["flare.struc.Structure", List["flare.struc.Structure"]]:
#     def __eq__(self, other):
#     def __init__(self):
#     def get_next_structure(self) -> Structure:
#     def write_file(self):
#     def __init__(self):
#     def get_next_force(self, *args, **kwargs) -> "np.ndarray":
#     def pre_force(self, *args, **kwargs):
#     def __init__(
#         self, frames: List[Structure] = None, iterate_strategy: Union[int, str] = 1
#     ):
#     def get_next_structure(self) -> Union[Structure, None]:
#     def cur_frame(self):
#     def cur_forces(self):
#     def get_next_force(self, index: int = -1) -> "np.ndarray":
#     def __getitem__(self, item):
#     def __len__(self):
#     def __next__(self):
#     def __iter__(self):
#     def append(self, frame: Structure):
# def get_unique_species(species: List[Any]) -> (List, List[int]):


# Path: tests/mgp_test.py
import os, shutil, re
import numpy as np
from flare import env, struc

        gp_force = np.sum(gp_f, axis=0)
        gp_energy = np.sum(gp_e, axis=0)
        print("gp_e", gp_e)
        print("gp_f")
        print(gp_f)

        map_ind = mgp_model.find_map_index(spc)
        xyzs = np.zeros_like(xyzs)
        xyzs[:, 0] = np.ones_like(xyzs[:, 0])
        f, _, _, e = mgp_model.maps[map_ind].predict(lengths, xyzs)

        assert np.allclose(gp_force, f, rtol=1e-2)
        assert np.allclose(gp_energy, e, rtol=1e-2)


def get_triplet_env(r1, r2, r12, grid_env):
    grid_env.bond_array_3 = np.array([[r1, 1, 0, 0], [r2, 0, 0, 0]])
    grid_env.cross_bond_dists = np.array([[0, r12], [r12, 0]])
    print(grid_env.ctype, grid_env.etypes)

    return grid_env


def get_grid_env(GP, species, bodies):
    if isinstance(GP.cutoffs, dict):
        max_cut = np.max(list(GP.cutoffs.values()))
    else:
        max_cut = np.max(GP.cutoffs)
    big_cell = np.eye(3) * 100
    positions = [[(i + 1) / (bodies + 1) * 0.1, 0, 0] for i in range(bodies)]

    grid_struc = struc.Structure(big_cell, species, positions)