# -*- coding: utf-8 -*-

def read_terminology_entry(terminology_entry):
    terminology_items = terminology_entry.split('~')

    terminology = {}
    terminology['contextName'] = terminology_items[0]

    terminology['category'] = terminology_items[1].split("^")
    terminology['type'] = terminology_items[2].split("^")
    typeModifier = terminology_items[3].split("^")
    if any(item!='' for item in typeModifier):
        terminology['typeModifier'] = typeModifier

    anatomicRegion = terminology_items[5].split("^")
    if any(item!='' for item in anatomicRegion):
        terminology['anatomicContextName'] = terminology_items[4]
        terminology['anatomicRegion'] = anatomicRegion
        anatomicRegionModifier = terminology_items[6].split("^")
        if any(item!='' for item in anatomicRegionModifier):
            terminology['anatomicRegionModifier'] = anatomicRegionModifier

    return terminology

def read_segmentation_info(filename):
    import nrrd
    header = nrrd.read_header(filename)
    segmentation_info = {}
    segments = []
    segment_index = 0
    while True:
        prefix = "Segment{0}_".format(segment_index)
        if not prefix + "ID" in header.keys():
            break
        segment = {}
        segment["index"] = segment_index
        segment["color"] = [float(i) for i in header[prefix + "Color"].split(" ")]  # Segment0_Color:=0.501961 0.682353 0.501961
        segment["colorAutoGenerated"] = int(header[prefix + "ColorAutoGenerated"]) != 0  # Segment0_ColorAutoGenerated:=1
        segment["extent"] = [int(i) for i in header[prefix + "Extent"].split(" ")]  # Segment0_Extent:=68 203 53 211 24 118
        segment["id"] = header[prefix + "ID"]  # Segment0_ID:=Segment_1
        segment["labelValue"] = int(header[prefix + "LabelValue"])  # Segment0_LabelValue:=1
        segment["layer"] = int(header[prefix + "Layer"])  # Segment0_Layer:=0
        segment["name"] = header[prefix + "Name"]  # Segment0_Name:=Segment_1
        segment["nameAutoGenerated"] = int(header[prefix + "NameAutoGenerated"]) != 0  # Segment0_NameAutoGenerated:=1
        # Segment0_Tags:=Segmentation.Status:inprogress|TerminologyEntry:Segmentation category and type - 3D Slicer General Anatomy list
        # ~SCT^85756007^Tissue~SCT^85756007^Tissue~^^~Anatomic codes - DICOM master list~^^~^^|
        tags = {}
        tags_str = header[prefix + "Tags"].split("|")
        for tag_str in tags_str:
            tag_str = tag_str.strip()
            if not tag_str:
                continue
            key, value = tag_str.split(":", maxsplit=1)
            if key=="TerminologyEntry":
                segment["terminology"] = read_terminology_entry(value)
            elif key=="Segmentation.Status":
                segment["status"] = value
            else:
                tags[key] = value
        segment["tags"] = tags
        segments.append(segment)
        segment_index += 1
    segmentation_info["segments"] = segments
    return segmentation_info

def segment_from_name(segmentation_info, segment_name):
    for segment in segmentation_info["segments"]:
        if segment_name == segment["name"]:
            return segment
    raise KeyError("segment not found by name " + segment_name)

def segment_names(segmentation_info):
    names = []
    for segment in segmentation_info["segments"]:
        names.append(segment["name"])
    return names

def extract_segments(voxels, header, segmentation_info, segment_names_to_label_values):
    import numpy as np
    # Create empty array from last 3 dimensions (output will be flattened to a 3D array)
    output_voxels = np.zeros(voxels.shape[-3:])
    # Copy non-segmentation fields to the extracted header
    output_header = {}
    for key in header.keys():
        if not re.match("^Segment[0-9]+_.+", key):
            output_header[key] = header[key]
    # Copy extracted segments
    dims = len(voxels.shape)
    for output_segment_index, segment_name_to_label_value in enumerate(segment_names_to_label_values):
        # Copy relabeled voxel data
        segment = segment_from_name(segmentation_info, segment_name_to_label_value[0])
        input_label_value = segment["labelValue"]
        output_label_value = segment_name_to_label_value[1]
        if dims == 3:
            output_voxels[voxels == input_label_value] = output_label_value
        elif dims == 4:
            inputLayer = segment["layer"]
            output_voxels[voxels[inputLayer,:,:,:] == input_label_value] = output_label_value
        else:
            raise ValueError("Voxel array dimension is invalid")
        # Copy all segment fields corresponding to this segment
        for key in header.keys():
            prefix = "Segment{0}_".format(segment["index"])
            matched = re.match("^" + prefix + "(.+)", key)
            if not matched:
                continue
            field_name = matched.groups()[0]
            if field_name == "LabelValue":
                value = output_label_value
            elif field_name == "Layer":
                # Output is a single layer (3D volume)
                value = 0
            else:
                value = header[key]
            output_header["Segment{0}_".format(output_segment_index) + field_name] = value
    # Remove unnecessary 4th dimension (volume is collapsed into 3D)
    if dims == 4:
        # Remove "none" from "none (0,1,0) (0,0,-1) (-1.2999954223632812,0,0)"
        output_header["space directions"] = output_header["space directions"][-3:,:]
        # Remove "list" from "list domain domain domain"
        output_header["kinds"] = output_header["kinds"][-3:]
    return output_voxels, output_header
