import importlib.resources

import numpy as np
import pandas as pd

__version__ = importlib.resources.read_text(__name__, '__version__')


def load_optotrak(file, delimiter='\t', max_rows=None) -> pd.DataFrame:
    """Load data from an Optotrak tab- or comma-separated values file.

    TSV files generated by the NDI software are a little weird, and need special
    handling.

    Parameters
    ----------
    file
        Filename.
    delimiter : str, optional
        Delimiter for the file. (default: '\t')
    max_rows : int, optional
        Maximum number of data rows to read. Use None to read all rows.
        (default: None)
    """
    #---------------------------------------------------------------------------
    # Position data
    #
    # Load the numeric data, replacing missing values with NaN. The final column
    # of data has an extraneous tab; this makes 'genfromtxt' think there's an
    # extra column that isn't there.
    #---------------------------------------------------------------------------
    data = np.genfromtxt(
        file,
        delimiter=delimiter,
        skip_header=5,
        max_rows=max_rows,
        missing_values=[''],
        filling_values=[float('nan')],
    )[:, :-1]

    frame = data[:, 0].astype('uint')
    x = data[:, 1::3]
    y = data[:, 2::3]
    z = data[:, 3::3]

    #---------------------------------------------------------------------------
    # Marker names
    #
    # Read marker names from the unique values of the header.
    #---------------------------------------------------------------------------
    header = pd.Series(
        np.loadtxt(
            file,
            delimiter=delimiter,
            skiprows=4,
            max_rows=1,
            dtype=object,
        )[1:])  # Drop 'Frame'
    # Split '<name> xyz', select '<name>' from resulting lists, then get uniques
    # '<name> xyz' -> ['<name>', 'xyz'] -> '<name>' -> unique
    names = header.str.split(' ').apply(lambda l: l[0]).unique()

    #---------------------------------------------------------------------------
    # Metadata
    #
    # Number of frames, sample rate, units
    #---------------------------------------------------------------------------
    metadata = dict(np.loadtxt(file, delimiter=': ', max_rows=3, dtype=object))

    data = pd.DataFrame(
        {
            'x': x.flat,
            'y': y.flat,
            'z': z.flat,
        },
        index=pd.MultiIndex.from_product([frame, names],
                                         names=['Frame', 'Name']),
    )
    data.attrs['Count'] = int(metadata['Number of frames'])
    data.attrs['Frequency'] = float(metadata['Frequency'])
    data.attrs['Units'] = metadata['Units']

    return data
