#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Aug 29 11:30:17 2019

@author: vand@dtu.dk
"""

import numpy as np
import scipy.io
import scipy.ndimage
import matplotlib.pyplot as plt
import st3d

# reading in the data
volume = scipy.io.loadmat('example_data_3D/multi_cube.mat')['vol']

# computing structure tensor and orientations
sigma = 0.5;
rho = 2;
S = st3d.structure_tensor(volume, sigma, rho)
val,vec = st3d.eig_special(S)
    
# visualization - use arrow keys to inspect the volume
st3d.show_vol_flow(volume, vec[0:2].reshape((2,)+volume.shape), s=10, double_arrow = True)
st3d.show_vol_orientation(volume, vec, coloring = st3d.fan_coloring)

    
#%% AND FOR A ROTATED SAMPLE...
volume = np.transpose(volume, (1, 0, 2))
S = st3d.structure_tensor(volume, sigma, rho)
val,vec = st3d.eig_special(S)
coloring_z = lambda v : st3d.fan_coloring(v[[2, 0, 1]]) # rotating coloring
st3d.show_vol_orientation(volume, vec, coloring = coloring_z)

#%% DISTANCE FROM A DIRECTION
u = np.array([0,0,1])
dist = st3d.tensor_vector_distance(S,u)
st3d.show_vol(dist.reshape(volume.shape))
 
# #%% FLOW MAKES SENSE ONLY WITH ROUGHLY UD FIBERS
# fxy = solve_flow_2D(S).reshape((2,)+volume.shape)

#%% K-MEANS LIKE CLUSTERING, INITILIZED BY 4 DIRECTIONS
t = np.sqrt(1/2)
u_clusters = np.array([[1,0,0],[0,0,1],[t,0,t],[-t,0,t]]).T # initalizing close to desired solution   
dist = st3d.tensor_vector_distance(S,u_clusters)
assignment = np.argmin(dist,axis = 1)
 
S_clusters = np.zeros((6,u_clusters.shape[1]))
for r in range(10): # iterations of k-means
    for i in range(u_clusters.shape[1]): # collecting ST for all voxels in a cluster
        S_clusters[:,i] = np.mean(S[:,assignment==i],axis=1)
    val,vec = st3d.eig_special(S_clusters) # estimating new cluster orientation
    print(f'Iter {r}: moved for {np.sum(u_clusters-vec)}')
    u_clusters = vec
    dist = st3d.tensor_vector_distance(S,u_clusters)
    assignment = np.argmin(dist,axis = 1)

st3d.show_vol(assignment.reshape(volume.shape), cmap='jet')