#!/usr/bin/env python
from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
import nibabel as nib
from nipype.interfaces import fsl
from nipype import Node, Workflow
from os.path import abspath
import subprocess as sub
import shlex as sh
import sys
import os

def removeFiles(paths):
    for element in paths:
        if os.path.exists(element):
            os.remove(element)
        else:
            print("No se puede remover el archivo %s" % element)

def write(v,id):
    fo = open("../web/app/public/densities/partial_density"+id+".txt",'w')
    fo.write(str(len(v))+",\n")
    fo.close() #La primera posicion debe ser la longitud del vector
    for i in range (0,len(v)):
        if(i == len(v)-1):#no le ponga la "," al dato final
            fo = open("../web/app/public/densities/partial_density"+id+".txt",'a')
            fo.write(str(v[i]))
            fo.close()
        else:
            fo = open("../web/app/public/densities/partial_density"+id+".txt",'a')
            fo.write(str(v[i])+",\n")
            fo.close()

def partial_density_weighted(data, data_mask, totalPixels):
    densities = []
    for i in range(0,data.shape[2]):
        xPixels = np.sum(data[:,:,i]) #suma de los pixeles en este corte
        wPixels = 0 if (xPixels == 0) else xPixels / totalPixels #peso total de ese corte en todo el volumen
        densities.insert(i,wPixels)
    return densities

def density(file_name, file_name_mask,id):
    img = nib.load(file_name)
    img_mask = nib.load(file_name_mask)
    data = img.get_data()
    data_mask = img_mask.get_data()
    num_data = data_mask.sum()
    sum_data = data.sum()
    density = (sum_data / num_data)
    print("Densidad Total : " , density)
    fo = open("../data/density"+id+".txt",'w')
    fo.write(str(density))
    fo.close()
    nib.aff2axcodes(img.affine)
    return partial_density_weighted(data, data_mask, num_data)

def threshold_CT(file_name,id):
    try:
        img = nib.load(file_name.decode("utf-8"))
        data = img.get_data()
        indices = data < 0
        data[indices] = 0
        indices = data > 100
        data[indices] = 0
        nib.save(img,'../data/threshold_brain'+id+'.nii.gz')
        return '../data/threshold_brain'+id+'.nii.gz'
    except:
        print("Error en los archivos ingresados")
        sys.exit(1)


def smooth_image(file_name):
    smooth = fsl.IsotropicSmooth(in_file=file_name, sigma=3, out_file=file_name)
    res = smooth.run()

def nifti_converter(output_dir, input_dir):
    command = 'dicomtonifti -brz -o '+output_dir+' '+input_dir;
    args = sh.split(command)
    print (command)
    p = sub.Popen(args, stdout=sub.PIPE, stderr=sub.PIPE)
    output, errors = p.communicate()
    print("ERRORES: ", errors)
    #sys.exit(1)
    output = output[:-1]
    return output
    #print (output)

def plot_slice(fname, z_idx=5):
	# Load the image and collect the data
	# and orientation information
	img = nib.load(fname)
	data = img.get_data()
	aff = img.get_affine()

	# Find the center of the brain matrix
	ctr = np.dot(np.linalg.inv(aff), [0, 0, 0, 1])[:3]
	#print(data[:,1,10])

	# Plot the data
	vmin, vmax = (0, 1) if data.dtype == np.int16 else (30, 150)
	plt.imshow(np.rot90(data[:, :, ctr[2] + z_idx]), cmap="gray", vmin=vmin, vmax=vmax)
	plt.gca().set_axis_off()

def skullStripping(file_name,id):
    skullstrip = fsl.BET()
    skullstrip.inputs.in_file = file_name
    skullstrip.inputs.out_file = "../data/T1_bet_nipype"+id+".nii.gz"
    skullstrip.inputs.mask = True
    skullstrip.inputs.frac = 0.01
    skullstrip.inputs.vertical_gradient = 0.0
    print(skullstrip.cmdline)
    skullstrip.run()
    return "../data/T1_bet_nipype"+id+".nii.gz","../data/T1_bet_nipype"+id+"_mask.nii.gz"


    #f = plt.figure()#figsize=(12,4))
    #f.add_subplot(1,2,1)
    #plot_slice(file_name)
    #plt.title("T1")
    #f.add_subplot(1,2,2)
    #plot_slice("../data/T1_bet_nipype.nii.gz")
    #plt.title("T1 Without Skull")
    #plt.show()


def drawImages():
    f = plt.figure(figsize=(12,4))
    for i, img in enumerate(["../data/T1.nii.gz","../data/smoothflow/smooth/T1_smooth.nii.gz","../data/smoothflow/mask/T1_smooth_masked.nii.gz"]):
        f.add_subplot(1,3,i+1)
        plot_slice(img)
    plt.show()


def createWorkFlow():
	in_file = abspath("../data/T1.nii.gz")
	skullstrip = Node(fsl.BET(in_file=in_file, mask=True), name ="skullstrip")
	smooth = Node(fsl.IsotropicSmooth(in_file=in_file, fwhm=3), name="smooth")
	mask = Node(fsl.ApplyMask(), name="mask")
	wf = Workflow(name="smoothflow")
	wf.base_dir = "/home/john/Documents/Projects/brainDensitometry/brain-densitometry/data/"
	wf.connect(skullstrip,"mask_file", mask, "mask_file")
	wf.connect([(smooth, mask, [("out_file","in_file")])])
	wf.write_graph("workflow_graph.dot")
	wf.run()


if len(sys.argv) >= 3:
    input_dicom_files = sys.argv[1]
    id = sys.argv[2]
    print("id : ", id)
else:
    print("Too few arguments")
    sys.exit(1)

paths = [] #ruta de archivos para remover

#print("antes : ", input_dicom_files)
input_dicom_files = input_dicom_files.replace(" ","\ ")
#print("despues : ", input_dicom_files)

file_name = nifti_converter('../data', input_dicom_files)
print("file_name1: " , file_name)
paths.append(file_name)

file_name = threshold_CT(file_name,id);
print("file_name2: " , file_name)
paths.append(file_name)

#smooth_image('../data/threshold_brain.nii.gz')
#threshold_CT('../data/threshold_brain.nii.gz')

file_name = skullStripping(file_name,id);
print("file_name3: " , file_name)
paths.append(file_name[0])
paths.append(file_name[1])

cuts = density(file_name[0],file_name[1],id)
print("Densidad Parcial : " , np.sum(cuts))
write(cuts,id)
removeFiles(paths)

#charting(cuts)
#createWorkFlow()
#drawImages()reateWorkFlow()
#drawImages()
