import os

import cv2

import math

import imutils

import numpy as np

import pandas as pd

import matplotlib.pyplot as plt

import config

def resize(image, 

           x=config.IMAGE_PXL_SIZE_X,

           y=config.IMAGE_PXL_SIZE_Y):

    if not len(image):

        return np.array([])

    return np.stack([cv2.resize(scan, (x, y)) 

                     for scan in image])

def normalize(image):

    image = ((image - config.MIN_BOUND) / 

             (config.MAX_BOUND - config.MIN_BOUND))

    image[image>1] = 1.

    image[image<0] = 0.

    return image

def chunks(l, n):

    """Yield successive n-sized chunks from l."""

    for i in range(0, len(l), n):

        yield l[i:i + n]

def mean(l):

    if len(l):

        return sum(l) / len(l)

    return np.full(l.shape, config.OUT_SCAN, l.dtype)

def get_mean_chunk_slices(slices):

    if len(slices) < config.SLICES:

        print("New slices are less then required after getting mean images, adding padding.")

        return trim_and_pad(np.array(slices), config.SLICES) 

    new_slices = []

    for slice_chunk in np.array_split(slices, config.SLICES):

        slice_chunk = list(map(mean, zip(*slice_chunk)))

        new_slices.append(slice_chunk)

    return np.stack(new_slices)

def read_csv(input_file):

    return pd.read_csv(input_file, index_col=0)

def read_csv_column(input_file, columns=[0]):

    return pd.read_csv(input_file, usecols=columns).values.flatten()

def store_to_csv(patients, labels, csv_file_path):

    index = pd.Index(data=patients, name=config.ID_COLUMN_NAME)

    df = pd.DataFrame(data={config.COLUMN_NAME: labels}, 

                      columns=[config.COLUMN_NAME],

                      index=index)

    df.to_csv(csv_file_path)

def trim_and_pad(patient_img, slice_count, normalize_pad=True):

    slices, size_x, size_y = patient_img.shape

    if slices == slice_count:

        return patient_img

    if slices > slice_count:

        return patient_img[:slice_count]

    padding = np.full((slice_count-slices, size_x, size_y), 

        config.OUT_SCAN, patient_img.dtype)

    if normalize_pad:

        padding = normalize(padding)

    return np.vstack([patient_img, padding])

def trim_pad_slices(scans, pad_with_existing=True,

                    padding_value=config.BACKGROUND):

    slices, x, y = scans.shape

    if slices == config.SLICES:

        return scans

    if slices < config.SLICES:

        pad = config.SLICES - slices

        new_scans = []

        if pad > slices:

            # Double the size, scans are already ordered by slice location

            for scan in scans:

                new_scans.append(scan) 

                new_scans.append(scan)

            del scans

            scans = new_scans

            pad = config.SLICES - len(scans)

        if pad_with_existing:

            padding = []

            for slice_chunk in np.array_split(scans, pad):

                padding.extend(slice_chunk)

                padding.append(slice_chunk[-1])

            #contains originals also, doubles the last slice in the chunk

            return np.stack(padding)

        else:

            padding = np.full((pad, x, y), padding_value, scans.dtype)

            return np.vstack([scans, padding])

    trim = slices - config.SLICES

    trimmed = []

    for slice_chunk in np.array_split(scans, trim):

        trimmed.append(slice_chunk[1:])

    return np.vstack(trimmed)

def count_background_rows(image, background=config.BACKGROUND):

    return np.sum(np.all(image == background, axis=1))

def remove_background_rows(image, background=config.BACKGROUND):

    return image[40:image.shape[0]-40, 20:image.shape[1]-20]

def rotate_scans(scans, angle=10):

    return np.stack([imutils.rotate(scan, angle) for scan in scans])

def remove_background_rows_3d(scans, background=config.BACKGROUND):

    transformed = []

    for scan in scans:

        removed = remove_background_rows(scan, background)

        tr_scan = cv2.resize(removed, 

            (config.IMAGE_PXL_SIZE_X, config.IMAGE_PXL_SIZE_Y))

        transformed.append(tr_scan)

    return np.stack(transformed)

def store_patient_image(image_dir, image, patient_id):

    """

    Serializes the patient image.

    Image is a 3D numpy array - array from patient slices.

    If not existing image_dir is created.

    """

    if not os.path.exists(image_dir):

        os.makedirs(image_dir)

    np.savez_compressed(os.path.join(image_dir, patient_id), image)

def load_patient_image(image_dir, patient_id):

    """

    Load the serialized patient image.

    Image is a 3D array - array of patient slices, metadata,

    contained in the dicom format, is removed.

    """

    if '.npz' not in patient_id:

        patient_id += '.npz'

    with np.load(os.path.join(image_dir, patient_id)) as data:

        return data['arr_0']