'''
This file is responsible for the generation of RNN input and output sequences.

Written by
  Alessandro Motta <alessandro.motta@brain.mpg.de>
'''
import os
import numpy as np
from skimage.color import rgb2gray
from skimage.io import imread

data_dir = os.path.join(os.path.dirname(__file__), "data")
file_names = list(filter(lambda n: n.lower()[-3:] in ['jpg', 'png'],
                         os.listdir(data_dir)))


def _load_grayscale_image(file_name) -> np.ndarray:
    color_image = imread(os.path.join(data_dir, file_name))
    gray_image = rgb2gray(color_image)
    return gray_image


def _load_trace_from_image(file_name):
    image = _load_grayscale_image(file_name)

    trace = np.argmin(image, axis=0)
    trace = np.asarray(trace, dtype=np.float32)

    trace -= np.linspace(trace[0], trace[-1], trace.size)
    trace /= 2 * trace.std()

    return trace


traces = list(map(_load_trace_from_image, file_names))
num_classes = len(traces)
assert num_classes == 2

# Amplitude of input signal
in_amplitude = 1.0
debug = False


def _build_sequence(sample_len, trace):
    x = trace.size / sample_len * np.arange(sample_len)
    seq = np.interp(x, np.arange(trace.size), trace)
    return seq


def generator(sample_len):
    sequences = list(map(lambda t: _build_sequence(sample_len, t), traces))
    if debug: import matplotlib.pyplot as plt

    while True:
        cur_seq_idx = np.random.randint(num_classes)
        cur_seq_sign = 2 * cur_seq_idx - 1
        cur_seq = sequences[cur_seq_idx]

        cur_begin = np.random.randint(sample_len)
        cur_len = sample_len - cur_begin

        cur_in_seq = np.zeros(sample_len, dtype=np.float32)
        cur_in_seq[cur_begin:] = in_amplitude * cur_seq_sign

        cur_out_seq = np.zeros(sample_len, dtype=np.float32)
        cur_out_seq[cur_begin:] = cur_seq[:cur_len]

        # Add channel dimension
        cur_in_seq = np.reshape(cur_in_seq, (-1, 1))
        cur_out_seq = np.reshape(cur_out_seq, (-1, 1))

        if debug:
            plt.figure()
            plt.plot(cur_in_seq)
            plt.plot(cur_out_seq)
            plt.show()

        yield cur_in_seq, cur_out_seq


