import os

import numpy as np

import math

import random

import cv2 as cv

import nibabel as nib

import torch

# in: volume path

# out: volume data in array

def readVol(volpath):

    return nib.load(volpath).get_data()

# in: volume array

# out: comprise to uint8, put 0 where number<0

def to_uint8(vol):

    vol=vol.astype(np.float)

    vol[vol<0]=0

    return ((vol-vol.min())*255.0/vol.max()).astype(np.uint8)

# in: volume array

# out: comprise to uint8, put 0 where number<800

def IR_to_uint8(vol):

    vol=vol.astype(np.float)

    vol[vol<0]=0

    return ((vol-800)*255.0/vol.max()).astype(np.uint8)

# in: volume array

# out: hist equalized volume arrray

def histeq(vol):

    for slice_index in range(vol.shape[2]):

        vol[:,:,slice_index]=cv.equalizeHist(vol[:,:,slice_index])

    return vol

# in: volume array

# out: preprocessed array

def preprocessed(vol):

    for slice_index in range(vol.shape[2]):

        cur_slice=vol[:,:,slice_index]

        sob_x=cv.Sobel(cur_slice,cv.CV_16S,1,0)

        sob_y=cv.Sobel(cur_slice,cv.CV_16S,0,1)

        absX=cv.convertScaleAbs(sob_x)

        absY=cv.convertScaleAbs(sob_y)

        sob=cv.addWeighted(absX,0.5,absY,0.5,0)

        dst=cur_slice+0.5*sob

        vol[:,:,slice_index]=dst

    return vol

# in: index of slice, stack number, slice number

# out: which slice should be stacked

def get_stackindex(slice_index, stack_num, slice_num):

    assert stack_num%2==1, 'stack numbers must be odd!'

    query_list=[0]*stack_num

    for stack_index in range(stack_num):

        query_list[stack_index]=(slice_index+(stack_index-int(stack_num/2)))%slice_num

    return query_list

# in: volume array, stack number

# out: stacked img in list

def get_stacked(vol,stack_num):

    stack_list=[]

    stacked_slice=np.zeros((vol.shape[0],vol.shape[1],stack_num),np.uint8)

    for slice_index in range(vol.shape[2]):

        query_list=get_stackindex(slice_index,stack_num,vol.shape[2])

        for index_query_list,query_list_content in enumerate(query_list):

            stacked_slice[:,:,index_query_list]=vol[:,:,query_list_content].transpose()

        stack_list.append(stacked_slice.copy())

    return stack_list

# in: stacked img, rotate angle

# out: rotated imgs

def rotate(stack_list,angle,interp):

    for stack_list_index,stacked in enumerate(stack_list):

        raws,cols=stacked.shape[0:2]

        M=cv.getRotationMatrix2D(((cols-1)/2.0,(raws-1)/2.0),angle,1)

        stack_list[stack_list_index]=cv.warpAffine(stacked,M,(cols,raws),flags=interp)

    return stack_list

# in: T1 volume, foreground threshold, margin pixel numbers

# out: which region should be cropped

def calc_crop_region(stack_list_T1,thre,pix):

    crop_region=[]

    for stack_list_index,stacked in enumerate(stack_list_T1):

        _,threimg=cv.threshold(stacked[:,:,int(stacked.shape[2]/2)].copy(),thre,255,cv.THRESH_TOZERO)

        pix_index=np.where(threimg>0)

        if not pix_index[0].size==0:

            y_min,y_max=min(pix_index[0]),max(pix_index[0])

            x_min,x_max=min(pix_index[1]),max(pix_index[1])

        else:

            y_min,y_max=pix,pix

            x_min,x_max=pix,pix

        y_min=(y_min<=pix)and(0)or(y_min)

        y_max=(y_max>=stacked.shape[0]-1-pix)and(stacked.shape[0]-1)or(y_max)

        x_min=(x_min<=pix)and(0)or(x_min)

        x_max=(x_max>=stacked.shape[1]-1-pix)and(stacked.shape[1]-1)or(x_max)

        crop_region.append([y_min,y_max,x_min,x_max])

    return crop_region

# in: crop region for each slice, how many slices in a stack

# out: max region in a stacked img

def calc_max_region_list(region_list,stack_num):

    max_region_list=[]

    for region_list_index in range(len(region_list)):

        y_min_list,y_max_list,x_min_list,x_max_list=[],[],[],[]

        for stack_index in range(stack_num):

            query_list=get_stackindex(region_list_index,stack_num,len(region_list))

            region=region_list[query_list[stack_index]]

            y_min_list.append(region[0])

            y_max_list.append(region[1])

            x_min_list.append(region[2])

            x_max_list.append(region[3])

        max_region_list.append([min(y_min_list),max(y_max_list),min(x_min_list),max(x_max_list)])

    return max_region_list

# in: size, devider

# out: padded size which can be devide by devider

def calc_ceil_pad(x,devider):

    return math.ceil(x/float(devider))*devider

# in: stack img list, maxed region list

# out: cropped img list

def crop(stack_list,region_list):

    cropped_list=[]

    for stack_list_index,stacked in enumerate(stack_list):

        y_min,y_max,x_min,x_max=region_list[stack_list_index]

        cropped=np.zeros((calc_ceil_pad(y_max-y_min,16),calc_ceil_pad(x_max-x_min,16),stacked.shape[2]),np.uint8)

        cropped[0:y_max-y_min,0:x_max-x_min,:]=stacked[y_min:y_max,x_min:x_max,:]

        cropped_list.append(cropped.copy())

    return cropped_list

# in: stack lbl list, dilate iteration

# out: stack edge list

def get_edge(stack_list,kernel_size=(3,3),sigmaX=0):

    edge_list=[]

    for stacked in stack_list:

        edges=np.zeros((stacked.shape[0],stacked.shape[1],stacked.shape[2]),np.uint8)

        for slice_index in range(stacked.shape[2]):

            edges[:,:,slice_index]=cv.Canny(stacked[:,:,slice_index],1,1)

            edges[:,:,slice_index]=cv.GaussianBlur(edges[:,:,slice_index],kernel_size,sigmaX)

        edge_list.append(edges)

    return edge_list

if __name__=='__main__':

    T1_path='../../data/training/1/pre/reg_T1.nii.gz'

    vol=to_uint8(readVol(T1_path))

    print(vol.shape)

    print('vol[100,100,20]= ', vol[100,100,20])

    histeqed=histeq(vol)

    print('vol[100,100,20]= ', vol[100,100,20])

    print('query list: ', get_stackindex(1,5,histeqed.shape[2]))

    stack_list=get_stacked(histeqed,5)

    print(len(stack_list))

    print(stack_list[0].shape)

    angle=random.uniform(-15,15)

    print('angle= ', angle)

    rotated=rotate(stack_list,angle)

    print(len(rotated))

    region=calc_crop_region(rotated,50,5)

    max_region=calc_max_region_list(region,5)

    print(region)

    print(max_region)

    cropped=crop(rotated,max_region)

    for i in range(48):

        print(cropped[i].shape)