import os
import sys

import cv2
import numpy as np

from ivus_segmentation.preprocessing import downsample

NUM_CLASSES = 3


def eprint(*args, verbose=True, **kwargs):
    kwargs.pop("verbose", None)
    if verbose:
        print(*args, file=sys.stderr, **kwargs)


def get_data(name, data_dir, cached=True, cache_dir="_cache", verbose=False):
    try:
        os.mkdir(cache_dir)
    except FileExistsError:
        pass

    files = [os.path.join(cache_dir, f"{name}_{path}")
             for path in ("images", "masks")]

    if cached:
        eprint("\x1b[34m•\x1b[39m Loading cached data...", verbose=verbose)
        try:
            arrs = []
            for file in files:
                eprint(f"  loading {file}", verbose=verbose)
                with open(file, "rb") as f:
                    arrs.append(np.load(f))
            return arrs

        except FileNotFoundError:
            eprint(
                "\x1b[31m!\x1b[39m Could not find cached files. Building dataset...."
            )

    x_train, y_train_labels = _load("A", data_location=data_dir)

    x_train_ds, y_train_label_ds = downsample(
        x_train, y_train_labels, NUM_CLASSES)
    # x_test_ds, y_test_label_ds = downsample(x_test, y_test_labels, NUM_CLASSES)

    images = x_train_ds[:, 2:5, :, :].transpose([0, 2, 3, 1]) / 255
    masks = y_train_label_ds

    # images_test = x_test_ds[:, 2:5, :, :].transpose([0, 2, 3, 1]) / 255
    # masks_test = y_test_label_ds

    out = [images.astype(np.float32), masks.astype(np.float32)]
    for file, arr in zip(files, out):
        with open(file, "wb") as f:
            np.save(f, arr)

    return out


def _read_labels(directory, verbose=0):
    """
    Reads the labels from a given directory

    returns:
        A dictionary containing:
            'lum': array of labels of 2D coordinates
            'med': array of labels of 2D coordinates
    """
    labels_lum = []
    labels_med = []

    for file_name in sorted(os.listdir(directory)):
        points = []
        with open(os.path.join(directory, file_name)) as f:
            for line in f:
                c0, c1 = line.split(",")
                points.append((float(c0), float(c1)))

        if file_name.startswith("lum"):
            labels_lum.append(points)
        elif file_name.startswith("med"):
            labels_med.append(points)

    # We will have imbalances since the number of coordinates is not fixed
    # To use numpy array, we will extend every list of labels repeating the last element
    labels_lum = np.array(labels_lum, dtype=object)
    if type(labels_lum[0]) == list:  # The dataset has no fixed number of rows or cols
        m = max([len(labels_lum[i]) for i in range(len(labels_lum))])
        for i in range(len(labels_lum)):
            labels_lum[i] = np.concatenate(
                (
                    labels_lum[i],
                    np.repeat([labels_lum[i][-1]], m -
                              len(labels_lum[i]), axis=0),
                )
            )

    labels_med = np.array(labels_med, dtype=object)
    if type(labels_med[0]) == list:  # The dataset has no fixed number of rows or cols
        m = max([len(labels_med[i]) for i in range(len(labels_med))])
        for i in range(len(labels_med)):
            labels_med[i] = np.concatenate(
                (
                    labels_med[i],
                    np.repeat([labels_med[i][-1]], m -
                              len(labels_med[i]), axis=0),
                )
            )

    # We have np.arrays inside np.arrays, we have to undo this
    labels = {
        "lum": np.array([l.tolist() for l in labels_lum]),
        "med": np.array([l.tolist() for l in labels_med]),
    }

    if verbose:
        print(
            f"Annotation for the inner wall (lumen contours): {labels['lum'].shape}")
        print(
            f"Annotation for the outer wall (media/adventitia contours): {labels['med'].shape}"
        )

    return labels


# Helper functions
def _load(letter="A", verbose=0, data_location="."):
    """
    Reads the training dataset A or B

    returns:
        Two tuples (x_train, y_train_label), (x_test, y_test_label):
            x_train: dataset of images for training
            y_train_label: a dictionary containing:
                'lum': array of labels of 2D coordinates
                'med': array of labels of 2D coordinates
            x_test: dataset of images for testing
            y_test_label: a list of 3 dictionaries containing:
                'lum': array of labels of 2D coordinates
                'med': array of labels of 2D coordinates
    """
    path = os.path.join(data_location, "Data_set_" + letter, "DCM")

    # dataset is N x 384 x 384 x 5

    files = sorted(os.listdir(path))
    N = len(files) // 5
    if letter == 'B':
        samples = np.empty((N, 5, 384, 384))
    elif letter == 'A':
        samples = np.empty((N, 5, 512, 512))
    else:
        raise NotImplementedError()

    for i in range(N):
        for j in range(5):
            file_name = files[i * 5 + j]
            img = cv2.imread(os.path.join(path, file_name),  # type: ignore
                             cv2.IMREAD_GRAYSCALE)  # type: ignore
            samples[i, j, :, :] = img
    if verbose:
        print(f"Dataset {letter} has shape: {samples.shape}")

    # Read the labels
    labels = []
    y_train_directory = os.path.join(data_location, "Data_set_" + letter)
    for dir_name in ["LABELS_obs2_v2"]:
        directory = os.path.join(y_train_directory, dir_name)
        labels.append(_read_labels(directory, verbose))

    # x_test = dataset[num_training_labels:, :, :, :]

    # # Read test labels
    # y_test_labels = []
    # y_test_directory = os.path.join(data_location, "Data_set_" + letter)
    # for dir_name in ["LABELS_obs2_v1"]:  # , "LABELS_obs2_v1", "LABELS_obs2_v2"]:
    #     directory = os.path.join(y_test_directory, dir_name)
    #     y_test_labels.append(_read_labels(directory, verbose))

    return (samples, labels)
