import cv2

import numpy as np

import os

import dicom

from pylab import *

from mpl_toolkits.mplot3d import Axes3D

slices=[]

def dicomRead(path):

    for dirName, subdirList, fileList in os.walk(path):

        for filename in fileList:

            if ".dcm" in filename.lower():

                slices.append(dicom.read_file(os.path.join(dirName, filename)))

    slices.sort(key=lambda x: int(x.ImagePositionPatient[2]))

    try:

        slice_thickness = np.abs(slices[0].ImagePositionPatient[2] - slices[1].ImagePositionPatient[2])

    except:

        slice_thickness = np.abs(slices[0].SliceLocation - slices[1].SliceLocation)

    for s in slices:

        s.SliceThickness = slice_thickness

    image = np.stack([s.pixel_array for s in slices])

    return image[0:255,:,:]

def floodfill(image, start_point, value):

    height, width = image.shape[:2]

    points = [start_point]

    flag = [[0 for j in range(width)] for i in range(height)]

    flag[start_point[0]][start_point[1]] = 1

    origin_value = image[start_point[0]][start_point[1]]

    while len(points) > 0:

        pt = points.pop(0)

        dx = [0, 1, 0, -1]

        dy = [1, 0, -1, 0]

        for x, y in zip(dx, dy):

            if (0 <= pt[0] + x < height and 0 <= pt[1] + y < width and

                        origin_value == image[pt[0] + x][pt[1] + y] and

                        flag[pt[0] + x][pt[1] + y] == 0):

                flag[pt[0] + x][pt[1] + y] = 1

                points.append((pt[0] + x, pt[1] + y))

        image[pt[0]][pt[1]] = value

    return image

def switch_pixels(image, origin_value, value):

    for i in range(image.shape[0]):

        for j in range(image.shape[1]):

            if image[i][j] == origin_value:

                image[i][j] = value

    return image

def morphology_open(image):

    # morphology open operation

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))

    image = cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel)

    return image

def segment(image):

    threashold=625

    mask = np.copy(image)

    mask = morphology_open(mask)

    ret, mask = cv2.threshold(mask, threashold, 255, cv2.THRESH_BINARY_INV)

    # set margin to black

    h, w = mask.shape[:2]

    for i in range(h):

        for j in range(w):

            if ((i == 0 or j == 0 or i == h - 1 or j == w - 1) and

                        mask[i][j] != 0):

                mask = floodfill(mask, (i, j), 0)

    # fill holes in middle

    mask = floodfill(mask, (0, 0), -1)

    mask = switch_pixels(mask, 0, 255)

    mask = switch_pixels(mask, -1, 0)

    return mask

def test3D(ConstPixelDims = None):

    xx, yy = np.meshgrid(np.linspace(0, 1, 512), np.linspace(0, 1, 512))

    X = xx

    Y = yy

    #Z = i

    ax2 = gca(projection='3d')

    off = 4000 / (ConstPixelDims[2] - 1)

    for i in range(ConstPixelDims[2]):  # ConstPixelDims[2]

        tempImage = ArrayDicom[:, :, i]

        tempImage = np.ma.masked_where(tempImage < 1200, tempImage)

        print("i : " + str(i))

        print("tempImage.shape[0] " + str(tempImage.shape[0]))

        print("tempImage.shape[1] " + str(tempImage.shape[1]))

        x, y = ogrid[0:tempImage.shape[0], 0:tempImage.shape[1]]

        ##print("x : " + str(x) + " :: y : " + str(y))

        Z = 10 * np.ones(X.shape)

        ax = gca(projection='3d')

        # ax.plot_surface(x, y, Z, rstride=5, cstride=5, facecolors=tempImage,cmap='gray' )

        ax2.contourf(X, Y, tempImage, zdir='z', offset=i * off, antialiased=True)

        # ax2.set_zlim((0.,1.))

    show()