import numpy as np

STATES   = ("AtEnd", "Backward", "Turn", "Forward", "Expect", "Lick")

PATTERNS = dict(AtEnd="<AtEnd> Backward {turn}",
                Backward="AtEnd <Backward> {turn}",
                Turn="Backward <{turn}> Forward", # note a difference from the previous section
                Forward="{turn} <Forward>", # note a difference from the previous section
                Expect="{turn} Forward <Expect>",
                Lick="{turn} Forward Expect <Lick>")

def collect_epochs_with_pattern(trials, pattern, property="right_whisker_angle_deg"):
    ret = []
    for trial in trials:
        # we collect traces within the matched periods, trial-by-trial
        traces = []
        for start, stop in trial.get_timeranges(pattern):
            if (stop - start) < 2:
                continue
            # frames start from 1, so we have to adjust when reading from tracking data
            epoch = np.array(trial.tracking[property][start-1:stop])
            traces.append(epoch)

        # ignore cases when the pattern did not match in this trial
        if len(traces) == 0:
            continue

        # otherwise: add with subject/session information
        ret.append(dict(session=(trial.subject, trial.session), traces=traces))
    return ret

def normalize_time(trace, t_out):
    """sample temporal normalization using a 1-D trace.
    `t_out` is a sequence of normalized timepoints, ranging [0,1].
    """
    t_in = np.arange(0, trace.size) / (trace.size - 1)
    return np.interp(t_out, t_in, trace)

def get_normalized_traces(epochs, bins=100):
    """returns a numpy.ndarray of time-normalized traces.
    the shape of the returned array is (bins, ntraces).

    `epochs` must be in the form as they are collected
    by `collect_epochs_with_pattern()`."""
    t_out = np.linspace(0, 1, bins, endpoint=True)
    traces = []
    for epoch in epochs:
        for trace in epoch["traces"]:
            traces.append(normalize_time(trace, t_out))
    return np.stack(traces, axis=-1)