from __future__ import print_function

import os
import numpy as np
import nibabel as nib
import skfuzzy as fuzz
from sklearn.cluster import KMeans
from multiprocessing import Pool, cpu_count


def create_dir(path):
    if not os.path.isdir(path):
        os.makedirs(path)
    return


def load_nii(path):
    nii = nib.load(path)
    return nii.get_data(), nii.get_affine()


def save_nii(data, path, affine):
    nib.save(nib.Nifti1Image(data, affine), path)
    return


def extract_features(data):
    x_idx, y_idx, z_idx = np.where(data > 0)
    features = []
    for x, y, z in zip(x_idx, y_idx, z_idx):
        features.append([data[x, y, z], x, y, z])
    return np.array(features)


def kmeans_cluster(data, n_clusters):
    features = extract_features(data)
    intensities = features[..., 0].reshape((-1, 1))
    kmeans_model = KMeans(n_clusters=n_clusters, init="k-means++",
                          precompute_distances=True, verbose=0,
                          random_state=7, n_jobs=1,
                          max_iter=1000, tol=1e-6).fit(intensities)

    labels = np.zeros(data.shape)
    for l, f in zip(kmeans_model.labels_, features):
        labels[int(f[1]), int(f[2]), int(f[3])] = l + 1

    return labels


def fuzzy_cmeans_cluster(data, n_clusters):
    features = extract_features(data)
    intensities = features[..., 0].reshape((1, -1))

    cntr, u, u0, d, jm, p, fpc = fuzz.cluster.cmeans(
        intensities, n_clusters, 2, error=1e-6,
        maxiter=1000, init=None, seed=7)
    labels_ = np.argmax(u, axis=0)

    labels = np.zeros(data.shape)
    for l, f in zip(labels_, features):
        labels[int(f[1]), int(f[2]), int(f[3])] = l + 1

    return labels


def get_target_label(labels, data):
    labels_set = np.unique(labels)
    mean_intensities = []
    for label in labels_set[1:]:
        label_data = data[np.where(labels == label)]
        mean_intensities.append(np.mean(label_data))
    target_intensity = np.median(mean_intensities)  # GM
    # target_intensity = np.max(mean_intensities)  # WM
    # target_intensity = np.min(mean_intensities)  # CSF
    target_label = mean_intensities.index(target_intensity) + 1
    return target_label


def unwarp_segment(arg, **kwarg):
    return segment(*arg, **kwarg)


def segment(src_path, dst_path, labels_path=None, method="km"):
    print("Segment on: ", src_path)
    try:
        data, affine = load_nii(src_path)
        n_clusters = 3

        if method == "km":
            # Method 1 - KMeans
            labels = kmeans_cluster(data, n_clusters)
        elif method == "fcm":
            # Method 2 - Fuzzy CMeans
            labels = fuzzy_cmeans_cluster(data, n_clusters)

        target = get_target_label(labels, data)
        gm_mask = np.copy(labels).astype(np.float32)
        gm_mask[np.where(gm_mask != target)] = 0.333
        gm_mask[np.where(gm_mask == target)] = 1.
        data = data.astype(np.float32)
        gm = np.multiply(data, gm_mask)
        save_nii(labels, labels_path, affine)
        save_nii(gm, dst_path, affine)
    except RuntimeError:
        print("\tFalid on: ", src_path)

    return


parent_dir = os.path.dirname(os.getcwd())
data_dir = os.path.join(parent_dir, "data")
data_src_dir = os.path.join(data_dir, "ADNIEnhance")
data_dst_dir = os.path.join(data_dir, "ADNIKMSegment")
data_labels = ["AD", "NC"]
create_dir(data_dst_dir)

data_src_paths, data_dst_paths, labels_paths = [], [], []
for label in data_labels:
    src_label_dir = os.path.join(data_src_dir, label)
    dst_label_dir = os.path.join(data_dst_dir, label)
    create_dir(dst_label_dir)
    for subject in os.listdir(src_label_dir):
        data_src_paths.append(os.path.join(src_label_dir, subject))
        subj_name = subject.split(".")[0]
        dst_subj_dir = os.path.join(dst_label_dir, subj_name)
        create_dir(dst_subj_dir)
        data_dst_paths.append(os.path.join(dst_subj_dir, subject))
        labels_paths.append(os.path.join(dst_subj_dir, subj_name + "_labels.nii.gz"))

method = "km"  # "fcm" or "km"

# Test
# segment(data_src_paths[0], data_dst_paths[0], labels_paths[0])

# Multi-processing
subj_num = len(data_src_paths)
paras = zip(data_src_paths, data_dst_paths,
            labels_paths, [method] * subj_num)
pool = Pool(processes=cpu_count())
pool.map(unwarp_segment, paras)