<commit_before>"""
LAMMPS EEX I/O
"""

import pandas as pd
import math
import re
import numpy as np

# Python 2/3 compat
try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO

import eex
import logging

# AMBER local imports
from . import amber_metadata as amd

logger = logging.getLogger(__name__)


def _parse_format(string):
    """
    Parses an AMBER style format string.

    Example:

    >>> string = "%FORMAT(10I8)"
    >>> _parse_format(string)
    [10, int, 8]
    """
    if "FORMAT" not in string:
        raise ValueError("AMBER: Did not understand format line '%s'." % string)

    pstring = string.replace("%FORMAT(", "").replace(")", "").strip()
    ret = [x for x in re.split('(\d+)', pstring) if x]

    if ret[1] == "I":
        if len(ret) != 3:
            raise ValueError("AMBER: Did not understand format line '%s'." % string)
        ret[1] = int
        ret[0] = int(ret[0])
        ret[2] = int(ret[2])
    elif ret[1] == "E":
        if len(ret) != 5:
            raise ValueError("AMBER: Did not understand format line '%s'." % string)
        ret[1] = float
        ret[0] = int(ret[0])
        ret[2] = int(ret[2])
        # The .8 is not interesting to us
    elif ret[1] == "a":
        if len(ret) != 3:
            raise ValueError("AMBER: Did not understand format line '%s'." % string)
        ret[1] = str
        ret[0] = int(ret[0])
        ret[2] = int(ret[2])
    else:
        raise ValueError("AMBER: Type symbol '%s' not understood from line '%s'." % (ret[1], string))

    return ret


def _data_flatten(data, column_name, category_index, df_index_name):
    # Reorganize the data 2D -> 1D packing
    flat_data = data.values.flatten()
    index = np.arange(category_index, flat_data.shape[0] + category_index)
    dl_col_name = column_name

    # Build and curate the data
    df = pd.DataFrame({df_index_name: index, dl_col_name: flat_data})
    df.dropna(axis=0, how="any", inplace=True)
    return df


def read_amber_file(dl, filename, blocksize=5000):

    ### First we register relevant parameters
    for k, v in amd.atom_data_units.items():
        dl.register_atom_units(k, v)

    for order, form_name, units in amd.register_forms:
        dl.register_functional_forms(order, form_name, utype=units)

    ### First we need to figure out system dimensions
    max_rows = 100  # How many lines do we attempt to search?
    header_data = eex.utility.read_lines(filename, max_rows)

    sizes_dict = {}
    ret_data = {}
    found_sizes = False
    found_version = False
    for num, line in enumerate(header_data):

        # First try to figure out the version
        if "VERSION" in line:
            sline = line.strip().split()
            if "VERSION_STAMP" == sline[1]:
                ret_data["VERSION"] = sline[3]
                if ret_data["VERSION"] != "V0001.000":
                    raise ValueError("AMBER Read: Did not recognize version '%s'." % ret_data["VERSION"])
                found_version = True
            else:
                raise ValueError("AMBER Read: Could not understand version line.")

        # Get size_information
        elif "FLAG POINTERS" in line:
            ncols, dtype, width = _parse_format(header_data[num + 1])

            parsed_sizes = []
            for shift in range(4):
                dline = header_data[num + 2 + shift]
                # Make sure the leading white space is not being cut off each line
                if shift < 3:
                    dline = dline.rjust(80)
                dline = [dtype(dline[i:(i + width)]) for i in range(0, len(dline), width)]
                parsed_sizes.extend(dline)

            if len(parsed_sizes) != len(amd.size_keys):
                raise IndexError("AMBER Read: The size of the FLAG POINTERS does not match the FLAG NAMES.")

            for k, v in zip(amd.size_keys, parsed_sizes):
                sizes_dict[k] = v

            found_sizes = True

    if not found_sizes:
        raise KeyError("AMBER Read: Did not find FLAG POINTERS data.")

    if not found_version:
        raise KeyError("AMBER Read: Did not find VERSION_STAMP data.")

    ### Iterate over the primary data portion of the object

    # Figure out the size of each label
    label_sizes = {}
    for k, v in amd.data_labels.items():
        if isinstance(v[0], int):
            label_sizes[k] = v[0]
        elif v[0] in list(sizes_dict):
            label_sizes[k] = sizes_dict[v[0]]
        else:
            # print("%30s %40s %d" % (k, v[0], int(eval(v[0], sizes_dict))))
            label_sizes[k] = int(eval(v[0], sizes_dict))

    # print(label_sizes)
    # Find the start
    current_data_category = None
    current_data_type = None
    file_handle = open(filename, "r")
    counter = 0
    for line in file_handle:
        if counter > 100:
            raise KeyError("AMBER Read: Could not find the line a data category in the first 100 lines.")
        matched, name = eex.utility.line_fuzzy_list(line, list(amd.data_labels))
        if matched:
            current_data_category = name
            current_data_type = _parse_format(next(file_handle))

            break
        counter += 1

    # Build any required temporaries
    _current_topology_indices = {
        "bonds": [3, 0, np.array([])],
        "angles": [4, 0, np.array([])],
        "dihedrals": [5, 0, np.array([])],
    }

    # Iterate over the file
    while True:

        # Type out the sizes and types
        nsize = label_sizes[current_data_category]
        nrows = int(math.ceil(nsize / float(current_data_type[0])))
        dtypes = [current_data_type[1]] * current_data_type[0]
        widths = [current_data_type[2]] * current_data_type[0]

        # Read in the data, in chunks
        remaining_read = nrows
        num_blocks = int(math.ceil(nrows / float(blocksize)))
        category_index = 0
        for block in range(num_blocks):

            # Figure out the size of the read
            read_size = blocksize
            if remaining_read < blocksize:
                read_size = remaining_read

            # read_fwf will push the file pointer *at least 4* so lets just read it in
            if read_size < 4:
                data = []
                for x in range(read_size):
                    data.append(next(file_handle))
                    tmp_handle = StringIO("".join(data))
            else:
                tmp_handle = file_handle

            # Read in the data
            data = pd.read_fwf(tmp_handle, nrows=read_size, widths=widths, dtypes=dtypes, header=None)

            # 1D atom properties
            if current_data_category in list(amd.atom_property_names):

                df = _data_flatten(data, amd.atom_property_names[current_data_category], category_index, "atom_index")
                category_index += df.shape[0]

                # Scale charge by constant used in AMBER to get units of e
                if current_data_category == "CHARGE":
                    df['charge'] = df['charge'] / 18.2223

                # Add the data to DL
                dl.add_atoms(df, by_value=True)

            elif current_data_category in amd.store_other:
                df = _data_flatten(data, current_data_category, category_index, "res_index")

                # Force residue pointer to be type int
                if current_data_category == "RESIDUE_POINTER":
                    df = df.astype(int)
                category_index += df.shape[0]

                # Add the data to DL
                dl.add_other(current_data_category, df)

            # Store bond, angle, dihedrals
            elif current_data_category in list(amd.topology_store_names):
                category = current_data_category.split("_")[0].lower()

                mod_size, current_size, remaining_data = _current_topology_indices[category]

                data = data.values.ravel()
                data = data[np.isfinite(data)]

                # Prepend remaining data
                if remaining_data.size:
                    data = np.hstack((remaining_data, data))

                # Get current remaining and
                remaining = data.size % mod_size
                if remaining:<commit_msg>Fixes up amber_read after units overhaul<commit_after>"""
LAMMPS EEX I/O
"""

import pandas as pd
import math
import re
import numpy as np

# Python 2/3 compat
try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO

import eex
import logging

# AMBER local imports
from . import amber_metadata as amd

logger = logging.getLogger(__name__)


def _parse_format(string):
    """
    Parses an AMBER style format string.

    Example:

    >>> string = "%FORMAT(10I8)"
    >>> _parse_format(string)
    [10, int, 8]
    """
    if "FORMAT" not in string:
        raise ValueError("AMBER: Did not understand format line '%s'." % string)

    pstring = string.replace("%FORMAT(", "").replace(")", "").strip()
    ret = [x for x in re.split('(\d+)', pstring) if x]

    if ret[1] == "I":
        if len(ret) != 3:
            raise ValueError("AMBER: Did not understand format line '%s'." % string)
        ret[1] = int
        ret[0] = int(ret[0])
        ret[2] = int(ret[2])
    elif ret[1] == "E":
        if len(ret) != 5:
            raise ValueError("AMBER: Did not understand format line '%s'." % string)
        ret[1] = float
        ret[0] = int(ret[0])
        ret[2] = int(ret[2])
        # The .8 is not interesting to us
    elif ret[1] == "a":
        if len(ret) != 3:
            raise ValueError("AMBER: Did not understand format line '%s'." % string)
        ret[1] = str
        ret[0] = int(ret[0])
        ret[2] = int(ret[2])
    else:
        raise ValueError("AMBER: Type symbol '%s' not understood from line '%s'." % (ret[1], string))

    return ret


def _data_flatten(data, column_name, category_index, df_index_name):
    # Reorganize the data 2D -> 1D packing
    flat_data = data.values.flatten()
    index = np.arange(category_index, flat_data.shape[0] + category_index)
    dl_col_name = column_name

    # Build and curate the data
    df = pd.DataFrame({df_index_name: index, dl_col_name: flat_data})
    df.dropna(axis=0, how="any", inplace=True)
    return df


def read_amber_file(dl, filename, blocksize=5000):

    ### First we register relevant parameters
    for order, form_name, units in amd.register_forms:
        dl.register_functional_forms(order, form_name, utype=units)

    ### First we need to figure out system dimensions
    max_rows = 100  # How many lines do we attempt to search?
    header_data = eex.utility.read_lines(filename, max_rows)

    sizes_dict = {}
    ret_data = {}
    found_sizes = False
    found_version = False
    for num, line in enumerate(header_data):

        # First try to figure out the version
        if "VERSION" in line:
            sline = line.strip().split()
            if "VERSION_STAMP" == sline[1]:
                ret_data["VERSION"] = sline[3]
                if ret_data["VERSION"] != "V0001.000":
                    raise ValueError("AMBER Read: Did not recognize version '%s'." % ret_data["VERSION"])
                found_version = True
            else:
                raise ValueError("AMBER Read: Could not understand version line.")

        # Get size_information
        elif "FLAG POINTERS" in line:
            ncols, dtype, width = _parse_format(header_data[num + 1])

            parsed_sizes = []
            for shift in range(4):
                dline = header_data[num + 2 + shift]
                # Make sure the leading white space is not being cut off each line
                if shift < 3:
                    dline = dline.rjust(80)
                dline = [dtype(dline[i:(i + width)]) for i in range(0, len(dline), width)]
                parsed_sizes.extend(dline)

            if len(parsed_sizes) != len(amd.size_keys):
                raise IndexError("AMBER Read: The size of the FLAG POINTERS does not match the FLAG NAMES.")

            for k, v in zip(amd.size_keys, parsed_sizes):
                sizes_dict[k] = v

            found_sizes = True

    if not found_sizes:
        raise KeyError("AMBER Read: Did not find FLAG POINTERS data.")

    if not found_version:
        raise KeyError("AMBER Read: Did not find VERSION_STAMP data.")

    ### Iterate over the primary data portion of the object

    # Figure out the size of each label
    label_sizes = {}
    for k, v in amd.data_labels.items():
        if isinstance(v[0], int):
            label_sizes[k] = v[0]
        elif v[0] in list(sizes_dict):
            label_sizes[k] = sizes_dict[v[0]]
        else:
            # print("%30s %40s %d" % (k, v[0], int(eval(v[0], sizes_dict))))
            label_sizes[k] = int(eval(v[0], sizes_dict))

    # print(label_sizes)
    # Find the start
    current_data_category = None
    current_data_type = None
    file_handle = open(filename, "r")
    counter = 0
    for line in file_handle:
        if counter > 100:
            raise KeyError("AMBER Read: Could not find the line a data category in the first 100 lines.")
        matched, name = eex.utility.line_fuzzy_list(line, list(amd.data_labels))
        if matched:
            current_data_category = name
            current_data_type = _parse_format(next(file_handle))

            break
        counter += 1

    # Build any required temporaries
    _current_topology_indices = {
        "bonds": [3, 0, np.array([])],
        "angles": [4, 0, np.array([])],
        "dihedrals": [5, 0, np.array([])],
    }

    # Iterate over the file
    while True:

        # Type out the sizes and types
        nsize = label_sizes[current_data_category]
        nrows = int(math.ceil(nsize / float(current_data_type[0])))
        dtypes = [current_data_type[1]] * current_data_type[0]
        widths = [current_data_type[2]] * current_data_type[0]

        # Read in the data, in chunks
        remaining_read = nrows
        num_blocks = int(math.ceil(nrows / float(blocksize)))
        category_index = 0
        for block in range(num_blocks):

            # Figure out the size of the read
            read_size = blocksize
            if remaining_read < blocksize:
                read_size = remaining_read

            # read_fwf will push the file pointer *at least 4* so lets just read it in
            if read_size < 4:
                data = []
                for x in range(read_size):
                    data.append(next(file_handle))
                    tmp_handle = StringIO("".join(data))
            else:
                tmp_handle = file_handle

            # Read in the data
            data = pd.read_fwf(tmp_handle, nrows=read_size, widths=widths, dtypes=dtypes, header=None)

            # 1D atom properties
            if current_data_category in list(amd.atom_property_names):

                df = _data_flatten(data, amd.atom_property_names[current_data_category], category_index, "atom_index")
                category_index += df.shape[0]

                # Add the data to DL
                dl.add_atoms(df, by_value=True, utype=amd.atom_data_units)

            elif current_data_category in amd.store_other:
                df = _data_flatten(data, current_data_category, category_index, "res_index")

                # Force residue pointer to be type int
                if current_data_category == "RESIDUE_POINTER":
                    df = df.astype(int)
                category_index += df.shape[0]

                # Add the data to DL
                dl.add_other(current_data_category, df)

            # Store bond, angle, dihedrals
            elif current_data_category in list(amd.topology_store_names):
                category = current_data_category.split("_")[0].lower()

                mod_size, current_size, remaining_data = _current_topology_indices[category]

                data = data.values.ravel()
                data = data[np.isfinite(data)]

                # Prepend remaining data
                if remaining_data.size:
                    data = np.hstack((remaining_data, data))

                # Get current remaining and
                remaining = data.size % mod_size
                if remaining:
{'fraction_non_alphanumeric': 0.10299105766265804, 'fraction_numerical': 0.012087573234659267, 'mean_word_length': 2.101759755164499, 'pattern_counts': {'":': 17, '<': 9, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2252409', 'n_tokens_mistral': 5060, 'n_tokens_neox': 4799, 'n_words': 1400}