import numpy as np
import networkx as nx
from tissue_classifier import TissueClassifier

class ChanVeseTissueClassifier(TissueClassifier):

    def calibration_slice_search(self,image):

        z_dim = image.shape[2]

        max_metric = 0.0
        max_zidx = 0

        for zidx in range(0,z_dim):
            z_slice = image[:,:,zidx]
            csf_total = np.sum(z_slice==1.0)
            gm_total = np.sum(z_slice==2.0)
            wm_total = np.sum(z_slice==3.0)
            brain_total = csf_total + gm_total + wm_total
            highest_total = np.max(np.asarray([csf_total,gm_total,wm_total]))
            if (highest_total==0.0):
                metric = 0.0
            else:
                metric = brain_total*(csf_total/highest_total + gm_total/highest_total + wm_total/highest_total)
            if (metric>max_metric):
                max_metric = metric
                max_zidx = zidx
        return max_zidx



    def compute_segmentation(self,g,image):
        cut_value, [Vs,Vt] = nx.minimum_cut(g,'s','t',capacity='capacity')
        x_dim = image.shape[0]
        y_dim = image.shape[1]

        labels = np.zeros(image.shape)

        label_dict = {'00':2,'01':3,'10':0,'11':1}

        for xidx in range(0,x_dim):
            for yidx in range(0,y_dim):
                vp1_str = "("+str(xidx)+","+str(yidx)+"),1"
                vp2_str = "("+str(xidx)+","+str(yidx)+"),2"

                if (vp1_str in Vs):
                    phi1 = 1
                else:
                    phi1 = 0

                if (vp2_str in Vs):
                    phi2 = 1
                else:
                    phi2 = 0

                label_key=str(phi1)+str(phi2)
                labels[xidx,yidx] = label_dict[label_key]

        return labels


    def chan_vese_segmentation(self,data,c0,v,calibration_slice):

        ctol = 0.1
        n_slices = self.data_dim[2]
        mid_point = int(n_slices/2)

        labels_3d = np.zeros(self.data_dim)

        c = c0
        c_converged = False
        while (c_converged == False):
            g = self.generate_graph(calibration_slice,c,v)
            labels = self.compute_segmentation(g,calibration_slice)
            labels_linear = labels.reshape(self.data_dim[0]*self.data_dim[1])
            image_linear = calibration_slice.reshape(self.data_dim[0]*self.data_dim[1])
            c_new = np.zeros(len(c))

            for cidx in range(0,len(c)):
                c_new[cidx] = np.mean(image_linear[labels_linear==cidx])

            if (np.sum(np.abs(c_new-c))<ctol):
                c_converged = True
            c = c_new
            print ("c: " + str(c))

        print ("Level set c has converged.")
        for zidx in range(0,n_slices):
            print ("Processing slice " + str(zidx) + " of " + str(n_slices))
            z_slice = data[:,:,zidx]
            g = self.generate_graph(z_slice,c,v)
            labels_3d[:,:,zidx] = self.compute_segmentation(g,z_slice)

        return labels_3d

    def generate_graph(self,image,c,v):

        g = nx.DiGraph()

        x_dim = image.shape[0]
        y_dim = image.shape[1]

        g.add_node('t',x=-1,y=-1,p=0,vld=0)
        g.add_node('s',x=-1,y=-1,p=0,vld=0)

        for xidx in range(0,x_dim):
            for yidx in range(0,y_dim):
                vp1_str = "("+str(xidx)+","+str(yidx)+"),1"
                vp2_str = "("+str(xidx)+","+str(yidx)+"),2"

                c1_sub = np.power(image[xidx,yidx]-c[0],2)
                c2_sub = np.power(image[xidx,yidx]-c[1],2)
                c3_sub = np.power(image[xidx,yidx]-c[2],2)
                c4_sub = np.power(image[xidx,yidx]-c[3],2)

                A_cap = np.max(c2_sub-c4_sub,0)
                C_cap = np.max(c4_sub-c2_sub,0)
                B_cap = np.max(c4_sub-c3_sub,0)
                D_cap = np.max(c3_sub-c4_sub,0)
                E_cap = c1_sub+c4_sub-c2_sub-c3_sub

                g.add_node(vp1_str,x=xidx,y=yidx,p=1,vld=1)
                g.add_node(vp2_str,x=xidx,y=yidx,p=2,vld=1)
                g.add_edge('s',vp1_str,capacity=C_cap)
                g.add_edge('s',vp2_str,capacity=D_cap)
                g.add_edge(vp1_str,'t',capacity=A_cap)
                g.add_edge(vp2_str,'t',capacity=B_cap)
                g.add_edge(vp1_str,vp2_str,capacity=E_cap)


        for xidx in range(0,x_dim):
            for yidx in range(0,y_dim):
                for xn in range(-1,2):
                    for yn in range(-1,2):
                        ynidx = yn+yidx
                        xnidx = xn+xidx

                        skip_condition = (yn==0 and xn==0) \
                        or (xnidx < 0 or xnidx >= x_dim) \
                        or (ynidx < 0 or ynidx >= y_dim)

                        if (skip_condition):
                            continue
                        if (xn==0 or yn==0):
                            wpq = v
                        else:
                            wpq = v/np.sqrt(2)
                        vp1_str = "("+str(xidx)+","+str(yidx)+"),1"
                        vp2_str = "("+str(xidx)+","+str(yidx)+"),2"
                        vq1_str = "("+str(xnidx)+","+str(ynidx)+"),1"
                        vq2_str = "("+str(xnidx)+","+str(ynidx)+"),2"
                        vpq1_str = vp1_str+"_"+vq1_str
                        vpq2_str = vp2_str+"_"+vq2_str
                        g.add_node(vpq1_str,x=-1,y=-1,p=0,vld=0)
                        g.add_node(vpq2_str,x=-1,y=-1,p=0,vld=0)
                        g.add_edge(vp1_str,vpq1_str,capacity=wpq)
                        g.add_edge(vpq1_str,vq1_str,capacity=float('inf'))
                        g.add_edge(vp2_str,vpq2_str,capacity=wpq)
                        g.add_edge(vpq2_str,vq2_str,capacity=float('inf'))
        print ("Graph generated")
        return g


    def estimate_labels(self,arg_dict):
        print ("Start: Chan-Vese")
        v = arg_dict['v']
        nonzero_data_idx = self.data!=0.0
        nonzero_data = self.data[nonzero_data_idx]
        labels_k, means_k, stds_k = self.k_means_clustering(nonzero_data,3)

        labeled_image_k = np.zeros(self.data_dim[0]*self.data_dim[1]*self.data_dim[2])
        labeled_image_k[nonzero_data_idx] = labels_k
        labeled_image_k = labeled_image_k.reshape(self.data_dim)

        calibration_slice_idx = self.calibration_slice_search(labeled_image_k)

        data = self.data.reshape(self.data_dim)
        calibration_slice = data[:,:,calibration_slice_idx]
        means = means_k
        stds = stds_k
        labels = self.chan_vese_segmentation(data,np.concatenate([np.zeros(1),means]),v,calibration_slice)

        self.statistics['means'] = means
        self.statistics['stds'] = stds
        self.labeled_image = labels
        print ("Complete: Chan-Vese")
