# <!----------------BEGIN-HEADER------------------------------------>
# ## FitSNAP3
# A Python Package For Training SNAP Interatomic Potentials for use in the LAMMPS molecular dynamics package
#
# _Copyright (2016) Sandia Corporation.
# Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
# the U.S. Government retains certain rights in this software.
# This software is distributed under the GNU General Public License_
# ##
#
# #### Original author:
#     Aidan P. Thompson, athomps (at) sandia (dot) gov (Sandia National Labs)
#     http://www.cs.sandia.gov/~athomps
#
# #### Key contributors (alphabetical):
#     Mary Alice Cusentino (Sandia National Labs)
#     Nicholas Lubbers (Los Alamos National Lab)
#     Maybe me ¯\_(ツ)_/¯
#     Adam Stephens (Sandia National Labs)
#     Mitchell Wood (Sandia National Labs)
#
# #### Additional authors (alphabetical):
#     Elizabeth Decolvenaere (D. E. Shaw Research)
#     Stan Moore (Sandia National Labs)
#     Steve Plimpton (Sandia National Labs)
#     Gary Saavedra (Sandia National Labs)
#     Peter Schultz (Sandia National Labs)
#     Laura Swiler (Sandia National Labs)
#
# <!-----------------END-HEADER------------------------------------->

from fitsnap3.io.input import config
from os import path, listdir, stat
import numpy as np
from random import shuffle
from fitsnap3.parallel_tools import pt
from fitsnap3.io.output import output
# from natsort import natsorted


def _float_to_int(a_float):
    if a_float == 0:
        return int(a_float)
    if a_float/int(a_float) != 1:
        raise ValueError("Training and Testing Size must be interpretable as integers")
    return int(a_float)


class Scraper:

    def __init__(self, name):
        self.name = name
        self.group_types = {}
        self.group_table = []
        self.files = {}
        self.configs = {}
        self.tests = None
        self.convert = {"Energy": 1.0, "Force": 1.0, "Stress": 1.0, "Distance": 1.0}
        self.data = {}

        self._init_units()

    def scrape_groups(self):
        group_dict = {k: config.sections["GROUPS"].group_types[i]
                      for i, k in enumerate(config.sections["GROUPS"].group_sections)}
        self.group_table = config.sections["GROUPS"].group_table
        size_type = None
        testing_size_type = None

        for key in self.group_table:
            bc_bool = False
            training_size = None
            if 'size' in self.group_table[key]:
                training_size = self.group_table[key]['size']
                bc_bool = True
                size_type = group_dict['size']
            if 'training_size' in self.group_table[key]:
                if training_size is not None:
                    raise ValueError("Do not set both size and training size")
                training_size = self.group_table[key]['training_size']
                size_type = group_dict['training_size']
            if 'testing_size' in self.group_table[key]:
                testing_size = self.group_table[key]['testing_size']
                testing_size_type = group_dict['testing_size']
            else:
                testing_size = 0
            if training_size is None:
                raise ValueError("Please set training size for {}".format(key))

            folder = path.join(config.sections["PATH"].datapath, key)
            folder_files = listdir(folder)
            for file_name in folder_files:
                if folder not in self.files:
                    self.files[folder] = []
                self.files[folder].append([folder + '/' + file_name, int(stat(folder + '/' + file_name).st_size)])
            shuffle(self.files[folder], pt.get_seed)
            nfiles = len(folder_files)
            if training_size < 1 or (training_size == 1 and size_type == float):
                if training_size == 1:
                    training_size = abs(training_size) * len(folder_files)
                else:
                    training_size = max(1, int(abs(training_size) * len(folder_files) - 0.5))
                if bc_bool and testing_size == 0:
                    testing_size = nfiles - training_size
            if testing_size != 0 and (testing_size < 1 or (testing_size == 1 and testing_size_type == float)):
                testing_size = max(1, int(abs(testing_size) * len(folder_files) - 0.5))
            training_size = _float_to_int(training_size)
            testing_size = _float_to_int(testing_size)
            if nfiles-testing_size-training_size < 0:
                raise ValueError("training size: {} + testing size: {} is greater than files in folder: {}".format(
                    training_size, testing_size, nfiles))
            output.screen(key, ": Detected ", nfiles, " fitting on ", training_size, " testing on ", testing_size)
            if self.tests is None:
                self.tests = {}
            self.tests[folder] = []
            for i in range(nfiles - training_size - testing_size):
                self.files[folder].pop()
            for i in range(testing_size):
                self.tests[folder].append(self.files[folder].pop())
            # self.files[folder] = natsorted(self.files[folder])

    # TODO : Fix divvy up to distribute groups evenly and based on memory
    def divvy_up_configs(self):
        groups = []
        group_list = []
        temp_list = []
        test_list = []
        for i, folder in enumerate(self.configs):
            for configuration in self.configs[folder]:
                if isinstance(configuration, list):
                    temp_list.append(configuration[0])
                else:
                    temp_list.append(configuration)
                groups.append(folder)

        self.configs = temp_list
        self.configs = pt.split_by_node(self.configs)

        if self.tests is not None:
            for i, folder in enumerate(self.tests):
                for file in self.tests[folder]:
                    test_list.append(file[0])
                    group_list.append(folder)
            test_list = pt.split_by_node(test_list)
            self.configs += test_list

        groups = pt.split_by_node(groups)
        group_list = pt.split_by_node(group_list)
        group_test = sorted(set(group_list))
        group_set = sorted(set(groups))
        group_counts = np.zeros((len(group_set) + len(group_test),), dtype='i')
        for i, group in enumerate(group_set+group_test):
            group_counts[i] = groups.count(group)

        pt.create_shared_array('files_per_group', len(self.configs), dtype='i')
        pt.shared_arrays['files_per_group'].array = group_counts

        pt.shared_arrays['files_per_group'].testing = 0
        if self.tests is not None:
            pt.shared_arrays['files_per_group'].testing = len(test_list)

        number_of_files_per_node = len(self.configs)
        pt.create_shared_array('number_of_atoms', number_of_files_per_node, dtype='i')
        pt.slice_array('number_of_atoms')
        self.configs = pt.split_within_node(self.configs)

    def scrape_configs(self):
        raise NotImplementedError("Call to virtual Scraper.scrape_configs method")

    def _init_units(self):
        if config.sections["REFERENCE"].units == "real":
            self.kb = 0.00198198665029335
        if config.sections["REFERENCE"].units == "metal":
            self.kb = 0.00008617333262145

    def _rotate_coords(self):
        # Transpose here because Lammps stores lattice vectors as columns,
        # QM stores lattice vectors as rows; After transposing lattice vectors are columns
        in_cell = np.asarray(self.data["QMLattice"]).T
        assert np.linalg.det(in_cell) > 0, "Input cell is not right-handed!"

        # Q matrix of QR decomposition is an orthogonal (rotation-like)
        # matrix whose inverse/transpose makes the input cell upper-diagonal:
        # input cell C = Q C';
        # runlammps-normalized cell C' = Q^T C.
        qmat, rmat = np.linalg.qr(in_cell)

        # Normalize signs of Q matrix to ensure positive diagonals of transformed cell;
        # QR decomposition algorithms don't always return a proper rotation
        ss = np.diagflat(np.sign(np.diag(rmat)))
        rot = ss @ qmat.T

        assert np.allclose(rot @ rot.T, np.eye(3)), "Rotation matrix not orthogonal"
        assert np.allclose(rot.T @ rot, np.eye(3)), "Rotation matrix not orthogonal"
        assert np.linalg.det(rot) > 0, "Rotation matrix is an improper rotation (det<0)"

        # ????
        # Cell transforms on first axis due to runlammps sotring lattice vectors as columns
        out_cell = rot @ in_cell

        # This assert is technically overkill, but checks that the new cell is right-handed
        assert np.linalg.det(out_cell) > 0, "New cell is not right-handed!"

        lower_triangle = out_cell[np.tril_indices(3, k=-1)]
        assert np.allclose(lower_triangle, 0, atol=1e-13), \
            f"Lower triangle of normalized cell has nonzero-elements: {lower_triangle}"

        # Positions and forces transform on the second axis
        # Stress transforms on both the first and second axis.
        self.data["Lattice"] = out_cell
        self.data["Positions"] = self.data["Positions"] * self.convert["Distance"] @ rot.T
        self.data["Forces"] = self.data["Forces"] * self.convert["Force"] @ rot.T
        self.data["Stress"] = rot @ (self.data["Stress"] * self.convert["Stress"]) @ rot.T
        self.data["Rotation"] = rot

    def _translate_coords(self):
        cell = self.data["Lattice"]
        position_in = self.data["Positions"]

        # Extra transposes because runlammps uses cells with latttice vectors as columns
        invcell = np.linalg.inv(cell.T).T
        # Fractional coordinates
        frac_coords = position_in @ invcell.T

        # Fix some rounding difficulties in divmod when within machine epsilon of zero
        frac_coords[np.isclose(frac_coords, 0, atol=1e-15)] = 0.

        trans_nums, cell_frac_coords = np.divmod(frac_coords, 1)

        assert (cell_frac_coords < 1).all(), "Fractional coordinates outside cell"
        assert (cell_frac_coords >= 0).all(), "Fractional coordinates outside cell"

        # If no translations are needed, return unmodified positions
        if (trans_nums == 0).all():
            self.data["Positions"] = position_in
            self.data["Translation"] = np.zeros_like(position_in, dtype=float)

        new_pos = cell_frac_coords @ cell.T
        trans_vec = trans_nums @ cell.T
        assert np.allclose(new_pos + trans_vec, position_in), "Translation failed to invert"
        self.data["Postions"] = new_pos
        self.data["Translation"] = trans_vec

    def _stress_conv(self, styles):

        if (config.sections["REFERENCE"].units == "metal" and
                list(styles["Stress"])[0] == "kbar" or list(styles["Stress"])[0] == "kB"):
            self.convert["Stress"] = 1000.0

    # def check_coords(self, cell, pos1, pos2):
    #     """Compares position 1 and position 2 with respect to periodic boundaries defined by cell"""
    #     invcell = np.linalg.inv(np.asarray(cell).T).T
    #
    #     # Fractional coordinates
    #     frac_1 = pos1 @ invcell.T
    #     frac_2 = pos2 @ invcell.T
    #     diff_frac = frac_2 - frac_1
    #
    #     # Assert that diff_frac is very close to an integer
    #     assert np.allclose(
    #         diff_frac,
    #         np.round(diff_frac),
    #         atol=1e-12, rtol=1e-12), "Coordinates are not close after shift." + \
    #                                  "Fractional coordinate Error:{}\nArray:\n{}".format(
    #                                      np.abs(diff_frac).max(), diff_frac)
    #     return True
    #
    # def check_volume(self, lattice, volume):
    #     assert np.allclose(np.linalg.det(lattice), volume, rtol=1e-10, atol=1e-10), \
    #         "Cell volume not equal to supplied volume!"
