"""
This file contains basic routines to load texture data. It is a
minimally modified version of Emmanuel's connectome/function/task/texture_data.py
"""
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_normalized(image_nr):
    gray_image = _load_grayscale_image(image_nr)
    gray_image -= gray_image.mean()
    gray_image /= gray_image.std()
    return gray_image


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


images = list(map(load_grayscale_image_normalized, range(len(file_names))))
num_classes = len(images)


def generator(sample_len, lims):
    lims = np.asarray(lims, dtype=np.float32)

    while True:
        cur_class = np.random.randint(num_classes)
        cur_sample = images[cur_class]
        cur_rows, cur_cols = cur_sample.shape

        cur_row = np.round(cur_rows * lims)
        cur_row = np.random.randint(*cur_row)
        cur_sample = cur_sample[cur_row, :].reshape(-1)

        cur_col = cur_cols - sample_len
        cur_col = np.random.randint(cur_col)
        cur_sample = cur_sample[cur_col:(cur_col + sample_len)]

        # NOTE(amotta): Flip sequence in 50 % of cases
        cur_do_flip = np.random.randint(2) == 0
        cur_sample = np.flip(cur_sample) if cur_do_flip else cur_sample

        # NOTE(amotta): Add "channels" dimension
        cur_sample = cur_sample.reshape((-1, 1))

        yield cur_sample, cur_class


def generators(sample_len, set_splits):
    set_splits = np.asarray(set_splits, dtype=np.float32)
    set_splits = np.concatenate(([0], set_splits.reshape(-1)))

    set_splits = set_splits.cumsum()
    set_splits = set_splits / set_splits[-1]

    lims = zip(set_splits[:-1], set_splits[1:])
    lims = tuple(map(lambda s: generator(sample_len, s), lims))
    return lims


def load_two_samples(sample_len):
    sample_one = load_grayscale_image_normalized(0)
    sample_one = sample_one[0, 0:sample_len].reshape((-1))

    sample_two = load_grayscale_image_normalized(1)
    sample_two = sample_two[0, 0:sample_len].reshape((-1))

    return sample_one, sample_two
