import numpy

from .columns import (
    META_DATA_EXPOSURE_ID,
    SETTINGS_EXPOSURE_TIME,
    META_DATA_PARAMETERS_TIME,
    SETTINGS_EXPOSURE_CHANNEL,
    META_DATA_PARAMETERS_CHANNEL,
)


def split_data_frame(data_frame, column):
    """ splits a data frame on unique column values """
    values = data_frame[column].unique()
    masks = {value: (data_frame[column] == value) for value in values}
    return {value: data_frame[mask] for value, mask in masks.items()}


def _set_exposure_data_from_parameters(data_frame):
    """infer the exposures from measurement parameters

    will raise a ValueError if the parameters contain NaNs
    """
    df = data_frame  # shorthand for cleaner code

    if (
        df[META_DATA_PARAMETERS_CHANNEL].hasnans
        or df[META_DATA_PARAMETERS_TIME].hasnans
    ):
        raise ValueError("Exposure Map: measurement parameters incomplete")

    df[SETTINGS_EXPOSURE_CHANNEL] = df[META_DATA_PARAMETERS_CHANNEL]
    df[SETTINGS_EXPOSURE_TIME] = df[META_DATA_PARAMETERS_TIME]
    return df


def apply_exposure_map(data_frame, exposure_map=None):
    """applies the parameters of a exposure map to the data frame

    exposure map:
        keys: must be the same as the exposure ids,
        values: objects with at least time and channel attributes

    if the exposure map is None, the values from the optionally parsed
    measurement parameters are used.

    will raise an ValueError, if the provided exposure map does not map to the
    exposure ids.
    """

    if exposure_map is None:
        return _set_exposure_data_from_parameters(data_frame)

    existing = set(data_frame[META_DATA_EXPOSURE_ID].unique())
    provided = set(exposure_map.keys())
    if existing != provided:
        raise ValueError(
            f"Exposure Map differs from data frame: {provided} != {existing}"
        )

    data_frame[SETTINGS_EXPOSURE_CHANNEL] = numpy.nan
    data_frame[SETTINGS_EXPOSURE_TIME] = numpy.nan
    for exposure_id, exposure_info in exposure_map.items():
        mask = data_frame[META_DATA_EXPOSURE_ID] == exposure_id
        data_frame.loc[mask, SETTINGS_EXPOSURE_CHANNEL] = exposure_info.channel
        data_frame.loc[mask, SETTINGS_EXPOSURE_TIME] = exposure_info.time
    return data_frame