import re

import pandas as pd

_CEA2034_COLUMNS = ['On Axis', 'Listening Window',
                    'Early Reflections', 'Sound Power', 'Sound Power DI', 'Early Reflections DI', 'DI offset']
_DIRECTIVITY_INDEX_COLUMNS = ['Sound Power DI', 'Early Reflections DI']
_EARLY_REFLECTIONS_COLUMNS = ['Floor Bounce',
                              'Ceiling Bounce', 'Front Wall Bounce', 'Side Wall Bounce', 'Rear Wall Bounce', 'Total Early Reflection']
_ESTIMATED_IN_ROOM_RESPONSE_COLUMNS = ['Estimated In-Room Response']
_HORIZONTAL_REFLECTIONS_COLUMNS = [
    'Front', 'Side', 'Rear', 'Total Horizontal Reflection']
_ANGLE_COLUMNS = (['On-Axis'] + [f'{angle}°' for absolute_angle in range(10, 180, 10)
                                 for angle in (absolute_angle, -absolute_angle)] + ['180°'])
_VERTICAL_REFLECTIONS_COLUMNS = [
    'Floor Reflection', 'Ceiling Reflection', 'Total Vertical Reflection']


_DIRECTIVITY_INDEX_COLUMNS = [(section, column) for section, columns in {
    'Directivity Index': _DIRECTIVITY_INDEX_COLUMNS,
}.items() for column in columns]
_SOUND_PESSURE_LEVEL_COLUMNS = [(section, column) for section, columns in {
    'CEA2034': _CEA2034_COLUMNS,
    'Early Reflections': _EARLY_REFLECTIONS_COLUMNS,
    'Estimated In-Room Response': _ESTIMATED_IN_ROOM_RESPONSE_COLUMNS,
    'Horizontal Reflections': _HORIZONTAL_REFLECTIONS_COLUMNS,
    'SPL Horizontal': _ANGLE_COLUMNS,
    'SPL Vertical': _ANGLE_COLUMNS,
    'Vertical Reflections': _VERTICAL_REFLECTIONS_COLUMNS,
}.items() for column in columns]

_COLUMNS = [(section,) + column for section, columns in {
    '[dB] Directivity Index ': _DIRECTIVITY_INDEX_COLUMNS,
    'Sound Pessure Level [dB]': _SOUND_PESSURE_LEVEL_COLUMNS,
}.items() for column in columns]
_COLUMNS_INDEX = columns = pd.MultiIndex.from_tuples(_COLUMNS)


def _fix_unnamed_columns(columns):
    # pd.read_table() expects the following multi-level column headers:
    #   A, A, A, A, B, B, B, B
    #   I, I, J, J, K, K, L, L
    #   X, Y, X, Y, X, Y, X, Y
    # But the data we have uses the following header format instead:
    #   A, B
    #   I, J, K, L
    #   X, Y, X, Y, X, Y, X, Y
    # When confronted with this header, pd.read_table() gets confused and
    # generates the following multi-level column index:
    #   A, _, _, _, B, _, _, _
    #   I, _, J, _, K, _, L, _
    #   X, Y, X, Y, X, Y, X, Y
    # Where "_" is some autogenerated column name in the form:
    # "Unnamed: 1_level_0"
    # This function restores the correct column names by replacing every
    # "Unnamed" column with the name of the last known column on that level.
    last_names = [None] * columns.nlevels

    def fix_column(column):
        for level, label in enumerate(column):
            if not label.startswith('Unnamed: '):
                last_names[level] = label
        return tuple(last_names)
    return pd.MultiIndex.from_tuples(fix_column(column) for column in columns.values)


def _index_by_frequency(data):
    # Expects input in the following form:
    #   (Additional top column levels)
    #   FR1                     FR2
    #   "Frequency [Hz]" value  "Frequency [Hz]" value
    #   42.42            1.234  42.42            2.345
    #   43.43            3.456  43.43            5.678
    # And reindexes it by the "Frequency [Hz]" column, producing:
    #          value
    #          (Additional top column labels)
    #          FR1    FR2
    #   42.42  1.234  2.345
    #   43.43  3.456  5.678
    preserve_column_level = list(range(data.columns.nlevels - 1))
    return (data
            # Move all columns levels except the bottommost one into the index
            .stack(level=preserve_column_level)
            # Drop the topmost (default) index level as it's not useful anymore
            .reset_index(level=0, drop=True)
            # Use the frequency as the new bottommost index level
            .set_index('Frequency [Hz]', append=True)
            # Move all other index levels back to columns
            .unstack(level=preserve_column_level))


def _cleanup_spl_column(column):
    # In "Sound Pessure Level [dB] / [2.83V 1m]", eliminates " / [2.83V 1m]", as
    # it varies between measurements
    match = re.match(r'^(Sound Pessure Level \[dB\])', column)
    return column if match is None else match.group(1)


def _load_fr(file):
    try:
        fr = pd.read_table(file, header=[0, 1, 2], thousands=',')
    except FileNotFoundError:
        # Some speakers are missing some files (e.g. a couple of early speakers
        # are missing "Directivity Index.txt"). We don't care about that here;
        # the notebook will validate the resulting set of curves after loading.
        return pd.DataFrame()
    fr.columns = _fix_unnamed_columns(fr.columns)
    return (fr
            .rename(columns=_cleanup_spl_column)
            .pipe(_index_by_frequency))


def load_speaker(dir):
    speaker = pd.concat(
        (_load_fr(dir / file) for file in [
            'CEA2034.txt', 'Directivity Index.txt', 'Early Reflections.txt',
            'Estimated In-Room Response.txt', 'Horizontal Reflections.txt',
            'SPL Horizontal.txt', 'SPL Vertical.txt',
            'Vertical Reflections.txt',
        ]),
        axis='columns')
    if speaker.isna().any(axis=None):
        # If this fires, it likely means something is wrong or corrupted in the
        # data files of the speaker (e.g. some frequencies present in some
        # columns/files but not others)
        raise AssertionError
    # Double-check that we don't have unexpected column labels.
    spurious_columns = speaker.columns.difference(_COLUMNS_INDEX)
    if not spurious_columns.empty:
        raise AssertionError(spurious_columns)
    # Some columns might be missing (e.g. a couple of early speakers are missing
    # the Directivity Index columns). We don't want downstream code to have to
    # deal with that edge case though, so fill any missing columns with N/A. The
    # downstream code will check for N/A and validate accordingly.
    return speaker.reindex(columns=_COLUMNS_INDEX)


def convert_angles(df):
    # Given a DataFrame with some of the columns in the following format:
    #   'On-Axis' '10°' '20°' '-10°' ...
    # Converts the above column labels to the following:
    #   0.0 10.0 20.0 -10.0
    def convert_label(label):
        if label == 'On-Axis':
            return 0.0
        stripped_label = label.strip('°')
        if stripped_label == label:
            return label
        try:
            return float(stripped_label)
        except ValueError:
            return label
    return df.rename(columns=convert_label)
