"""
Python wrapper for the gromacs library for their file formats.

Currently this supports xtc and tpr files within the `open` function.
Trajectories in trr format may be read with `.gromacs.reader.TRRReader`, which is experimental.
"""

import os

from .tpxio import TPXReader
from .xtcio import XTCReader
from .enxio import EDRFile
from .errors import FileTypeError
from .gromacs.reader import index_filename_for_xtc
from .gromacs.xtcindex import index_xtcfile

FILE_EXTENSIONS = {
    'xtc': XTCReader,
    'tpr': TPXReader,
    'edr': EDRFile,
}


def open(filename):
    """Open a supported gromacs file. Currently supported file formats: tpr, xtc."""
    ext = filename.split('.')[-1]
    if ext in FILE_EXTENSIONS:
        if ext in ['xtc']:
            indexfile = index_filename_for_xtc(filename)
            if not os.path.exists(indexfile):
                print('Generating Index for xtc file. This may take a while...')
                index_xtcfile(filename)
            return FILE_EXTENSIONS[ext](filename, indexfile)
        else:
            return FILE_EXTENSIONS[ext](filename)
    else:
        raise FileTypeError('Filetype {} not supported by pygmx.'.format(ext))