"""
Code for merging and munging trajectories from FAH datasets.

"""
##############################################################################
# imports
##############################################################################

from __future__ import print_function, division
import os, os.path
import glob
import tarfile
from mdtraj.formats.hdf5 import HDF5TrajectoryFile
import mdtraj as md
import tables
from mdtraj.utils.contextmanagers import enter_temp_directory
from mdtraj.utils import six
from natsort import natsorted
import time

##############################################################################
# globals
##############################################################################


def strip_water(allatom_filename, protein_filename, protein_atom_indices, min_num_frames=1):
    """Strip water (or other) atoms from a Core17, Core18, or OCore FAH HDF5 trajectory.

    Parameters
    ----------
    allatom_filename : str
        Path to HDF5 trajectory with all atoms.  This trajectory must have been generated by
        concatenate_core17 or concatenate_siegetank--e.g. it must include
        extra metadata that lists the XTC files (bzipped or in OCore directories) that
        have already been processed.  This file will not be modified.
    protein_filename : str
        Path to HDF5 trajectory with all just protein atoms.  This trajectory must have been generated by
        concatenate_core17 or concatenate_siegetank--e.g. it must include
        extra metadata that lists the XTC files (bzipped or in OCore directories) that
        have already been processed.  This file will be appended to.
    protein_atom_indices : np.ndarray, dtype='int'
        List of atom indices to extract from allatom HDF5 file.
    min_num_frames : int, optional, default=1
        Skip if below this number.

    """
    # Check integrity of trajectory if it exists.
    delete_trajectory_if_broken(allatom_filename)

    if not os.path.exists(allatom_filename):
        print("Skipping, %s not found" % allatom_filename)
        return

    trj_allatom = HDF5TrajectoryFile(allatom_filename, mode='r')

    print('all-atom trajectory %s has %d frames' % (allatom_filename, len(trj_allatom)))
    if len(trj_allatom) < min_num_frames:
        print("Must have at least %d frames in %s to proceed!" % (min_num_frames, allatom_filename))
        del trj_allatom
        return

    if hasattr(trj_allatom.root, "processed_filenames"):
        key = "processed_filenames"  # Core17, Core18 style data
    elif hasattr(trj_allatom.root, "processed_directories"):
        key = "processed_directories"  # Siegetank style data
    else:
        raise(ValueError("Can't find processed files in %s" % allatom_filename))

    # Check integrity of trajectory if it exists.
    delete_trajectory_if_broken(protein_filename)

    # Open the stripped trajectory.
    trj_protein = HDF5TrajectoryFile(protein_filename, mode='a')

    try:
        trj_protein._create_earray(where='/', name=key, atom=tables.StringAtom(1024), shape=(0,))
        trj_protein.topology = trj_allatom.topology.subset(protein_atom_indices)
    except tables.NodeError:
        pass

    n_frames_allatom = len(trj_allatom)
    try:
        n_frames_protein = len(trj_protein)
    except tables.NoSuchNodeError:
        n_frames_protein = 0

    filenames_allatom = getattr(trj_allatom.root, key)
    filenames_protein = getattr(trj_protein._handle.root, key)  # Hacky workaround of MDTraj bug #588

    n_files_allatom = len(filenames_allatom)
    n_files_protein = len(filenames_protein)
    print("Found %d,%d filenames and %d,%d frames in %s and %s, respectively." % (n_files_allatom, n_files_protein, n_frames_allatom, n_frames_protein, allatom_filename, protein_filename))

    if n_frames_protein > n_frames_allatom:
        raise(ValueError("Found more frames in protein trajectory (%d) than allatom trajectory (%d)" % (n_frames_protein, n_frames_allatom)))

    if n_files_protein > n_files_allatom:
        raise(ValueError("Found more filenames in protein trajectory (%d) than allatom trajectory (%d)" % (n_files_protein, n_files_allatom)))

    if n_frames_protein == n_frames_allatom or n_files_allatom == n_files_protein:
        if not (n_frames_protein == n_frames_allatom and n_files_allatom == n_files_protein):
            raise(ValueError("The trajectories must match in BOTH n_frames and n_filenames or NEITHER."))
        else:
            print("Same number of frames and filenames found, skipping.")
            del trj_allatom, trj_protein
            return

    trj_allatom.seek(n_frames_protein)  # Jump forward past what we've already stripped.
    coordinates, time, cell_lengths, cell_angles, velocities, kineticEnergy, potentialEnergy, temperature, alchemicalLambda = trj_allatom.read()
    trj_protein.write(coordinates=coordinates[:, protein_atom_indices], time=time, cell_lengths=cell_lengths, cell_angles=cell_angles)  # Ignoring the other fields for now, TODO.

    filenames_protein.append(filenames_allatom[n_files_protein:])
    del trj_allatom, trj_protein

def delete_trajectory_if_broken(filename, verbose=True):
    """
    Check the integrity of an MDTraj trajectory, deleting it if it is broken.

    Parameters
    ----------
    filename : str
       The trajectory filename.
    verbose : bool, optional, default=True
       If True, write some logging messages if broken trajectories are detected.

    """
    if os.path.exists(filename):
        try:
            trj = md.open(filename)
        except Exception as e:
            msg = "The integrity of trajectory file '%s' was compromised; deleting so that it will be regenerated.\n" % filename
            msg += "\n"
            msg += str(e)
            if verbose:
                print(msg)
            os.path.unlink(filename)

        # Clean up.
        del trj

def concatenate_core17(path, top_filename, output_filename, maxtime=None, maxpackets=None):
    """Concatenate tar bzipped XTC files created by Folding@Home Core17.
    This version accepts only filenames and paths.

    Parameters
    ----------
    path : str
        Path to directory containing "results-*.tar.bz2".  E.g. a single CLONE directory.
    top_filename : str
        Filepath to read Topology for system
    output_filename : str
        Filename of output HDF5 file to generate.
    maxpackets : int, optional, default=None
        If specified, will stop processing after `maxpackets` results packets have been processed
    maxtime : int, optional, default=None
        If specified, will stop processing after `maxtime` seconds have passed.

    Notes
    -----
    We use HDF5 because it provides an easy way to store the metadata associated
    with which files have already been processed.
    """

    # Open topology file.
    top = md.load(top_filename % vars())

    # Glob file paths and return result files in sequential order.
    glob_input = os.path.join(path, "results-*.tar.bz2")
    filenames = glob.glob(glob_input)
    filenames = natsorted(filenames)

    print("Concatenating XTC files from '%s' into '%s' [%d results packets found]" % (path, output_filename, len(filenames)))

    # If no result files are present, return.
    if len(filenames) <= 0:
        del top
        return

    # Check integrity of trajectory if it exists.
    delete_trajectory_if_broken(output_filename)

    # Open trajectory for appending.
    trj_file = HDF5TrajectoryFile(output_filename, mode='a')

    MAX_FILEPATH_LENGTH = 1024 # Is this large enough?
    try:
        # TODO: Store MD5 hashes instead of filenames?
        trj_file._create_earray(where='/', name='processed_filenames',atom=trj_file.tables.StringAtom(MAX_FILEPATH_LENGTH), shape=(0,))
        trj_file.topology = top.topology
    except trj_file.tables.NodeError:
        # Object already exists; skip ahead.
        pass

    result_packets_processed = 0
    initial_time = time.time()
    try:
        for filename in filenames:
            # Check that we haven't violated our filename length assumption
            if len(filename) > MAX_FILEPATH_LENGTH:
                msg = "Filename is longer than hard-coded MAX_FILEPATH_LENGTH limit (%d > %d). Increase MAX_FILEPATH_LENGTH and rebuild." % (len(filename), MAX_FILEPATH_LENGTH)
                print(msg)
                raise Exception(msg)
            # Check if we have already processed this file
            if six.b(filename) in trj_file._handle.root.processed_filenames:  # On Py3, the pytables list of filenames has type byte (e.g. b"hey"), so we need to deal with this via six.
                print("Already processed %s" % filename)
                continue
            # Extract frames from trajectory in a temporary directory
            absfilename = os.path.abspath(filename)
            with enter_temp_directory():
                # Extract frames
                archive = tarfile.open(absfilename, mode='r:bz2')
                archive.extract("positions.xtc")
                trj = md.load("positions.xtc", top=top)
                print("   appending %d frames from '%s' to '%s'" % (trj.n_frames, filename, output_filename))
                for frame in trj:
                    trj_file.write(coordinates=frame.xyz, cell_lengths=frame.unitcell_lengths, cell_angles=frame.unitcell_angles, time=frame.time)
                os.unlink("positions.xtc")
                del archive, trj

                # Append list of processed files
                trj_file._handle.root.processed_filenames.append([filename])

                # Flush data
                trj_file.flush()

            # Track statistics on processed packets
            elapsed_time = time.time() - initial_time
            result_packets_processed += 1

            # Return if we have processed the requested number of results packets.
            if maxpackets and (result_packets_processed >= maxpackets):
                break

    except RuntimeError:
        print("Cannot munge %s due to damaged XTC %s or mismatch with topology file." % (path, filename))

    # Clean up.
    trj_file.close()
    del top, trj_file

def concatenate_ocore(path, top_filename, output_filename):
    """Concatenate XTC files created by Siegetank OCore.

    Parameters
    ----------
    path : str
        Path to stream directory containing frame directories /0, /1, /2
        etc.
    top_filename : str
        Filepath to read Topology for system
    output_filename : str
        Filename of output HDF5 file to generate.

    Notes
    -----
    We use HDF5 because it provides an easy way to store the metadata associated
    with which files have already been processed.
    """
    # Open topology file.
    top = md.load(top_filename % vars())

    sorted_folders = sorted(os.listdir(path), key=lambda value: int(value))
    sorted_folders = [os.path.join(path, folder) for folder in sorted_folders]

    if len(sorted_folders) <= 0:
        return

    trj_file = HDF5TrajectoryFile(output_filename, mode='a')

    try:
        trj_file._create_earray(where='/', name='processed_folders',atom=trj_file.tables.StringAtom(1024), shape=(0,))
        trj_file.topology = top.topology
    except trj_file.tables.NodeError:
        pass

    for folder in sorted_folders:
        if six.b(folder) in trj_file._handle.root.processed_folders:  # On Py3, the pytables list of filenames has type byte (e.g. b"hey"), so we need to deal with this via six.
            print("Already processed %s" % folder)
            continue
        print("Processing %s" % folder)
        xtc_filename = os.path.join(folder, "frames.xtc")
        trj = md.load(xtc_filename, top=top)

        for frame in trj:
            trj_file.write(coordinates=frame.xyz, cell_lengths=frame.unitcell_lengths, cell_angles=frame.unitcell_angles, time=frame.time)

        trj_file._handle.root.processed_folders.append([folder])
