# (c) 2020-2021, Luca Widmer  @  ETH Zurich
# Computer-assisted Applications in Medicine (CAiM) Group,  Prof. Orcun Goksel

import h5py
import sys
import os
import numpy as np
from skimage import io
import json
import site

site.addsitedir(os.getcwd())
from data_read.datasets import dtfm
from config import settings
from fiji_funcs import CoordinateUtils


"""
This script implements the conversion of Z-stacks (generated by FIJI) to a MINTIF HDF5 File
"""


def extract_patch_from_img(array, patch_index, patch_size, z_offset=0, mean=None, std=None):
    """
    adapted from old get_index_for_patch
	This is a function for getting the patch. This function is written because
	there can be some negative indices, for example patch indices can be [-20, 30, -20, 30, -1, 2].
	In this case, only the positive indices are used and the negative indices are padded with zeros.

	Patch extracted from whole image or annotation masks, zero-padded if necessary..
    @param std:
    @param mean:
    @param z_offset: odset from 0 in z direction. used to handle z stacks that dont contain whole image.
    @param array: Whole image or annotation masks where the patch will be extracted from.
    @param patch_index: Patch indices that should be extracted.
    @param patch_size: dimensions of patch in pixel
    @return: extracted patch
    """
    patch_index[0] -= z_offset
    patch_index[1] -= z_offset

    z, x, y = array.shape
    ww = [patch_size[0], patch_size[1], patch_size[2]]

    ret = np.zeros(ww)
    temp_patch_index = np.array(patch_index).copy()
    ww = [0, patch_size[0], 0, patch_size[1], 0, patch_size[2]]

    # if patch overlaps image boundry (needs 0 padding) offset image index
    if temp_patch_index[0] < 0:
        ww[0] -= temp_patch_index[0]
        temp_patch_index[0] = 0
    if temp_patch_index[2] < 0:
        ww[2] -= temp_patch_index[2]
        temp_patch_index[2] = 0
    if temp_patch_index[4] < 0:
        ww[4] -= temp_patch_index[4]
        temp_patch_index[4] = 0

    if temp_patch_index[1] > z:
        ww[1] -= temp_patch_index[1] - z
        temp_patch_index[1] = z
    if temp_patch_index[3] > x:
        ww[3] -= temp_patch_index[3] - x
        temp_patch_index[3] = x
    if temp_patch_index[5] > y:
        ww[5] -= temp_patch_index[5] - y
        temp_patch_index[5] = y
    if temp_patch_index[0] >= temp_patch_index[1]:
        temp_patch_index[0] = temp_patch_index[1] - 1

    insert = array[temp_patch_index[0]:temp_patch_index[1],
             temp_patch_index[2]:temp_patch_index[3],
             temp_patch_index[4]:temp_patch_index[5]]

    # normalize patch
    if not (mean is None or std is None):
        insert = np.divide(insert - mean, std)

    ret[ww[0]:ww[1], ww[2]:ww[3], ww[4]:ww[5]] = insert

    return ret


def get_overlaping_points(patch_ind, coords_voxel):
    """
    find points lying in this patch and return point coordinates converted to Patch origin (0,0,0 at first pixel of patch)
    ATTENTION: indices & Points should be ZXY
    """
    res_points = np.empty((0, 3), int)
    for point in coords_voxel:
        # check if ZXY point is within patch indices ZXY
        if patch_ind[0] <= point[0] <= patch_ind[1] and patch_ind[2] <= point[1] <= patch_ind[3] and patch_ind[4] <= \
                point[2] <= patch_ind[5]:
            # subtract index value of patch to get origin at first pixel of patch
            point_patch_ori = np.subtract(point, np.take(patch_ind, [0, 2, 4]))
            res_points = np.append(res_points, [point_patch_ori], axis=0)

    return res_points

# ---------------------------
# PARSE COMMANDLINE ARGS
# ---------------------------

# try:
images = sys.argv[1]
hdf5_file_path = sys.argv[2]
model_file = sys.argv[3]
compress = sys.argv[4] or True
indices_file = sys.argv[5]
metadata_file = sys.argv[6]

f = h5py.File(hdf5_file_path, "a")

print("add {} to {}".format(images, hdf5_file_path))
# find metadata file
image_name = os.path.basename(images).split('.tiff')[0]
# load metadata
with open(metadata_file, 'r') as m:
    metadata = json.load(m)
# load model params and fill in missing values with defaults
msettings = settings.get_settings(model_file)

patch_size = np.array(msettings['patch_size'])
desired_voxel_size = msettings['voxel_size']
out_patch_size=msettings["patch_size_out"]


if len(desired_voxel_size) == 2:
    desired_voxel_size = np.append(1, desired_voxel_size)
    print("converted desired voxel size to 3D: {}".format(desired_voxel_size))

if len(patch_size) == 2:
    patch_size = [1, patch_size[0], patch_size[1]]
    print("converted patch size to 3D: {}".format(patch_size))
if len (out_patch_size) ==2:
    out_patch_size = [1]+ out_patch_size

# parse metadata
orig_voxel_dimensions = metadata["voxel_dimensions"]
user_metadata = metadata["user_metadata"]
channel_names = metadata["channel_names"]
train_set_2D = metadata["training_data"] and metadata["annotations"] == "2D"
channel_is_label = ["label" in i for i in channel_names]
channel_is_marker = ["in_" in i for i in channel_names]
Coordinate_Dict = metadata["Coodinate_Channels"] or {}
channel_stats = metadata["channel_statistics"]

# parse index files
with open(indices_file, 'r') as i:
    indices = json.load(i)
    in_patch_ind = indices["in_patch_ind"]
    out_patch_ind = indices["out_patch_ind"]
    z_stacks = indices["z_stack"]
    z_stacks_orig = indices["z_stack_original"]
    z_top=z_stacks_orig[-1][1]

# initialize 0 list to count label pixel for every label channel and total_cnumber
cnumber = np.array([0] * (len(channel_names) + 1))
# dimensions: z, c, y, x
z_dim, _, x_dim, y_dim = metadata["image_dimensions"]
c_dim = len(channel_names)
print("image dimensions :\nz: {}, c: {}, x: {}, y: {}".format(z_dim, c_dim, x_dim, y_dim))

# find slices which contain annotations for all labels (only 2D trainig data)
all_labels_in_slice_contain_anot = [True] * z_dim
if train_set_2D:
    for k, i in channel_stats.items():
        if "label_" in k:
            all_labels_in_slice_contain_anot = [x and not (y) for x, y in
                                                zip(all_labels_in_slice_contain_anot, i["empty"])]

# sanity check
assert len(in_patch_ind) == len(out_patch_ind)

print("Sample_{}".format(len(f.keys())))
sample = f.create_group("Sample_{}".format(len(f.keys())))

# write metadata of sample and channels to sample (1 based index for channels)
sample.attrs["image_shape"] = (z_dim, c_dim, x_dim, y_dim)
sample.attrs["name"] = metadata["original_file_name"]
sample.attrs["voxel_size"] = desired_voxel_size
sample.attrs["dataset_type"] = msettings["dataset_type"]

for name, sta in channel_stats.items():
    if "in_channel_" not in name:
        continue
    sample.attrs["mean_{}".format(name)] = round(sta["mean"],4)
    sample.attrs["std_{}".format(name)] = round(sta["std"],4)


# ---------------------------
# WRITE PATCHES
# ---------------------------
for channel_n, channel_name in enumerate(channel_names):
    channel_index = channel_name[-1]
    patch_counter=0
    for patch_n, (in_patch, out_patch) in enumerate(zip(in_patch_ind, out_patch_ind)):
        # # report progress
        # if (len(in_patch_ind)>=200 and patch_n % 100 == 0) or len(in_patch_ind)<200 and patch_n % 10 == 0:
        #     print("Patch {} of {}, Channel {} of {}".format(patch_n, len(in_patch_ind), channel_n, c_dim))

        # skip slice if its 2D training dataset and there are no annotations in slice
        if train_set_2D and not all_labels_in_slice_contain_anot[in_patch[0]]:
            continue
        else:
            patch_name = "Patch_{}".format(patch_counter)
            patch_counter+=1
            # create new patch if it not already exists
            try:
                patch = sample[patch_name]
            except:
                patch = sample.create_group(patch_name)

        if channel_n == 0:
            patch.attrs["do_pad"] = msettings["padding"] != "valid"
            patch.attrs["inds_channel"] = in_patch
            patch.attrs["inds_label"] = out_patch_ind[patch_n]
            patch.attrs["overlap_patch"] = indices["overlap_patch"]

        # load only the necessary z-stack to read current patch,
        # the loop should be structured in a way that ensures no unnecessary (re-)loadign happens
        old_image = image_name
        image_name = "image#{}#z{}_{}.tiff".format(channel_name, in_patch[0], in_patch[1])
        # load new image here if necessary
        if old_image != image_name:
            print("load {}".format(image_name))
            stack_path = os.path.join(images, image_name)
            image = io.imread(stack_path).astype("float32")

        # assure image is 3D
        if len(image.shape) == 2:
            image = np.reshape(image, (1, image.shape[0], image.shape[1]))
        # if not all slices contain annotation in 2D test dataset, skip
        if train_set_2D and not np.sum(image, axis=(1, 2)).any():
            continue

        if channel_is_label[channel_n]:
            # extract patch without padding, no normalization for label channels,
            # should z_offset not always be identical to lower patch boundary?
            channel_image = extract_patch_from_img(image, out_patch, out_patch_size, z_offset=max(0, in_patch[0]))
            # save number of pixels in labels to cnumber
            cnumber[channel_n] += np.count_nonzero(channel_image > 0)
            cnumber[-1] += channel_image.shape[1] * channel_image.shape[2]

        # MARKER CHANNEL
        # do gaussian samplewise normalization with pre-calcualted mean & std
        elif channel_is_marker[channel_n]:
            if msettings['normalize_data']:
                channel_mean = channel_stats[channel_name]["mean"]
                channel_std = channel_stats[channel_name]["std"]
            else:
                channel_mean = None
                channel_std = None
            # extract patch with padding and normalize, z offset is the lower boundary of the currently loaded slice
            channel_image = extract_patch_from_img(image, in_patch.copy(), patch_size, z_offset=max(0, in_patch[0]),
                                                   mean=channel_mean, std=channel_std)
        # if result contains only 1 slice, save it as 2D image
        if channel_image.shape[0] == 1:
            channel_image = channel_image[0,...]
        channel = patch.create_dataset(channel_names[channel_n], shape=channel_image.shape, chunks=True,
                                       data=channel_image, compression="gzip" if compress else None,
                                       compression_opts=9)
        # write channel metadata
        channel.attrs["Information"] = metadata["user_metadata"][channel_n]
        # todo implement name
        channel.attrs["name"] = ""
        # todo implement element size
        channel.attrs["element_size"] = []

# writre cnumber (used in segmentation)
for i, cn in enumerate(cnumber):
    if i == len(cnumber) - 1:
        sample.attrs["total_cnumber"] = cn
    elif channel_is_label[i]:
        sample.attrs["{}_cnumber".format(channel_names[i])] = cn

# ---------------------------
# WRITE 3D COORDINATES
# ---------------------------
"""
After writing file, loop trough all patches, check if inds_channel overlap with point coordinates, 
if yes, append to table of this patch.
First convert coords to Patch coordinate (Origin is First Patch pixel)
"""
# only safe coordinates, if file is provided and is training data
if len(Coordinate_Dict.keys()) > 0 and metadata["training_data"]:
    try:
        sample.attrs["radius_um"] = msettings["cells_radius"]
        sample.attrs["radius_pixel"] = msettings["cells_radius_pixel"]
        # read coordinate file
        for channel in Coordinate_Dict.values():
            # coordinates passed by mask channel will already have this key and do not need to be translated
            if "coords_voxel" not in channel:
                # original_voxel_size, coords = FuncsUtils.read_coord_file(channel["FilePath"], convert_coords=False)
                # if statement is legacy code
                # if ("convert_coords") not in channel:
                #     convert_coords = True
                # else:
                #     convert_coords = channel["convert_coords"]
                coords, original_voxel_size = CoordinateUtils.get_scaled_voxel_coords(channel["FilePath"], desired_voxel_size)
                channel["coords_voxel"], original_voxel_size = coords, original_voxel_size


    except Exception as e:
        raise UserWarning("Error while parsing Coordinate File:\n{}".format(e))

    for patchname, patch in sample.items():
        for channel_dict in Coordinate_Dict.values():
            patch_ind = patch.attrs["inds_label"]
            # get points in this patch
            patch_points = get_overlaping_points(patch_ind, channel_dict["coords_voxel"])
            # write data
            channel_name = "label_coordinates_{}".format(int(channel_dict["label"]))
            channel = patch.create_dataset(channel_name, data=patch_points, compression="gzip" if compress else None)
            # write patch metadata
            patch.attrs["do_pad"] = msettings["padding"] != "valid"
            channel.attrs["inds_channel"] = patch.attrs["inds_channel"]
            channel.attrs["inds_label"] = patch.attrs["inds_label"]
            channel.attrs["name"] = ""
            # todo implement radius_voxels
            channel.attrs['radius_um'] = ""
            channel.attrs['radius_voxels'] =""

    f.close()
    print("Convert Coordinates to Density Map")
    # write Density Map channel
    df = dtfm.DTFM(hdf5_file_path, msettings)
    df.coordinates2dm()
f.close()
print("done")
