#!/usr/bin/env python
# coding: utf-8
from __future__ import print_function

"""
CUDA_VISIBLE_DEVICES=1 python -W ignore Genesis_Chest_CT.py \
--note genesis_chest_ct \
--arch Vnet \
--input_rows 64 \
--input_cols 64 \
--input_deps 32 \
--nb_class 1 \
--verbose 1 \
--batch_size 16 \
--scale 32 \
--data generated_cubes
"""
import warnings
warnings.filterwarnings('ignore')
import os
# import keras
# print("Keras = {}".format(keras.__version__))
# import tensorflow as tf


import copy
import sys
import math
import random
import shutil
import numpy as np

from tqdm import tqdm
from scipy.special import comb
from sklearn import metrics
# from keras.callbacks import LambdaCallback, TensorBoard
from skimage.transform import resize
from optparse import OptionParser


def bernstein_poly(i, n, t):
    """
     The Bernstein polynomial of n, i as a function of t
    """

    return comb(n, i) * ( t**(n-i) ) * (1 - t)**i

def bezier_curve(points, nTimes=1000):
    """
       Given a set of control points, return the
       bezier curve defined by the control points.

       Control points should be a list of lists, or list of tuples
       such as [ [1,1], 
                 [2,3], 
                 [4,5], ..[Xn, Yn] ]
        nTimes is the number of time steps, defaults to 1000

        See http://processingjs.nihongoresources.com/bezierinfo/
    """

    nPoints = len(points)
    xPoints = np.array([p[0] for p in points])
    yPoints = np.array([p[1] for p in points])

    t = np.linspace(0.0, 1.0, nTimes)

    polynomial_array = np.array([ bernstein_poly(i, nPoints-1, t) for i in range(0, nPoints)   ])
    
    xvals = np.dot(xPoints, polynomial_array)
    yvals = np.dot(yPoints, polynomial_array)

    return xvals, yvals

def data_augmentation(inputs, prob=0.5):
    # augmentation by flipping
    cnt = 3
    while random.random() < prob and cnt > 0:
        degree = random.choice([0, 1, 2])
        for idx in range(len(inputs)):
            inputs[idx] = np.flip(inputs[idx],axis=degree)
        cnt = cnt - 1

    return inputs

def nonlinear_transformation(x, prob=0.5):
    if random.random() >= prob:
        return x
    points = [[0, 0], [random.random(), random.random()], [random.random(), random.random()], [1, 1]]
    xpoints = [p[0] for p in points]
    ypoints = [p[1] for p in points]
    xvals, yvals = bezier_curve(points, nTimes=100000)
    if random.random() < 0.5:
        # Half change to get flip
        xvals = np.sort(xvals)
    else:
        xvals, yvals = np.sort(xvals), np.sort(yvals)
    nonlinear_x = np.interp(x, xvals, yvals)
    return nonlinear_x

def local_pixel_shuffling(x, prob=0.5,num_block=500):
    if random.random() >= prob:
        return x
    flag = True
    if len(x.shape)==3:
        x = x[np.newaxis,:,:,:]
        flag = False
        
    image_temp = copy.deepcopy(x)
    orig_image = copy.deepcopy(x)
    _, img_rows, img_cols, img_deps = x.shape

    for _ in range(num_block):
        block_noise_size_x = random.randint(1, 3)
        block_noise_size_y = random.randint(1, img_cols//8)
        block_noise_size_z = random.randint(1, img_deps//8)

        noise_x = random.randint(0, img_rows-block_noise_size_x)
        noise_y = random.randint(0, img_cols-block_noise_size_y)
        noise_z = random.randint(0, img_deps-block_noise_size_z)
        window = orig_image[0, noise_x:noise_x+block_noise_size_x, 
                               noise_y:noise_y+block_noise_size_y, 
                               noise_z:noise_z+block_noise_size_z,
                           ]
        window = window.flatten()
        np.random.shuffle(window)
        window = window.reshape((block_noise_size_x, 
                                 block_noise_size_y, 
                                 block_noise_size_z))
        image_temp[0, noise_x:noise_x+block_noise_size_x, 
                      noise_y:noise_y+block_noise_size_y, 
                      noise_z:noise_z+block_noise_size_z] = window
    local_shuffling_x = image_temp
    if not flag:
        local_shuffling_x = np.squeeze(local_shuffling_x)
    return local_shuffling_x

def image_in_painting(x,min_size=10,max_size=20):
    flag = True
    if len(x.shape)==3:
        x = x[np.newaxis,:,:,:]
        flag = False
        
    _, img_rows, img_cols, img_deps = x.shape
    block_noise_size_x = random.randint(min_size, max_size)
    block_noise_size_y = random.randint(min_size, max_size)
    block_noise_size_z = random.randint(min_size, max_size)
    noise_x = random.randint(3, img_rows-block_noise_size_x-3)
    noise_y = random.randint(3, img_cols-block_noise_size_y-3)
    noise_z = random.randint(3, img_deps-block_noise_size_z-3)
    result = copy.deepcopy(x)
    result[:, 
      noise_x:noise_x+block_noise_size_x, 
      noise_y:noise_y+block_noise_size_y, 
      noise_z:noise_z+block_noise_size_z] = random.random()
    if not flag:
        result = np.squeeze(result)
    return result


def image_out_painting(x,min_size_z=10,min_size_x=10,max_size_z=10,max_size_x=10):
    flag = True
    if len(x.shape)==3:
        x = x[np.newaxis,:,:,:]
        flag = False
    
    min_val,max_val = np.amin(x),np.amax(x)
    val_range = max_val - min_val
    margin_size = 2
    _, img_rows, img_cols, img_deps = x.shape
#     print (min_size_x,max_size_x,img_rows-margin_size*2)
    block_noise_size_x = random.randint(min_size_z, min(max_size_z,img_rows-margin_size*2))
    block_noise_size_y = random.randint(min_size_x, min(max_size_x,img_cols-margin_size*2))
    block_noise_size_z = random.randint(min_size_x, min(max_size_x,img_deps-margin_size*2))
#     print ('block_noise_size_x',block_noise_size_x,block_noise_size_y,block_noise_size_z)
    noise_x = random.randint(margin_size, img_rows-block_noise_size_x-margin_size)
    noise_y = random.randint(margin_size, img_cols-block_noise_size_y-margin_size)
    noise_z = random.randint(margin_size, min(img_deps-block_noise_size_z-margin_size,int(1.5*margin_size)))
    image_temp = copy.deepcopy(x)
    result = copy.deepcopy(x)
#     result = np.random.rand(x.shape[0], x.shape[1], x.shape[2], x.shape[3], ) * 1.0
    result = np.random.rand(x.shape[0], x.shape[1], x.shape[2], x.shape[3], ) * val_range + min_val
#     print ('start x and end x',noise_x,noise_x+block_noise_size_x)
#     print ('start y and end y',noise_y,noise_y+block_noise_size_y)
    
    result[:, 
      noise_x:noise_x+block_noise_size_x, 
      noise_y:noise_y+block_noise_size_y, 
      noise_z:noise_z+block_noise_size_z] = image_temp[:, noise_x:noise_x+block_noise_size_x, 
                                                          noise_y:noise_y+block_noise_size_y, 
                                                        noise_z:noise_z+block_noise_size_z]
    if not flag:
        result = np.squeeze(result)
    return result
                
def random_crop(image,input_shape,deltas,aug):
    
    margin_z,margin_y,margin_x = [x/2 for x in input_shape]
    centers = [x/2 for x in image.shape]
    if not aug:
        deltas = [0,0,0]
    result = image[int(centers[0]+deltas[0]-margin_z):int(centers[0]+deltas[0]+margin_z),
                  int(centers[1]+deltas[1]-margin_y):int(centers[1]+deltas[1]+margin_y),
                  int(centers[2]+deltas[2]-margin_x):int(centers[2]+deltas[2]+margin_x)]
#     print ('after cropping result shape is ',result.shape,image.shape)
    return result
    

def roll_data(data):
    roller = np.round(float(data.shape[0]/7))
    oz,oy,ox = np.random.randint(-roller, roller+1, 3)

    result_data = np.roll(np.roll(np.roll(data, ox, 2), oy, 1),oz,0)
    return result_data
