import extension as ext_
import plot as ot_
import soma as soma_

import matplotlib.pyplot as pl_
import numpy as np_
import skimage.color as cl_
import skimage.io as io_
import skimage.measure as ms_
import skimage.morphology as mp_
import time as tm_


print(f"STARTED: {tm_.strftime('%a, %b %d %Y @ %H:%M:%S')}\n")
start_time = tm_.time()


soma_t = soma_.soma_t
extension_t = ext_.extension_t


# --- Parameters
run = ("soma", "extension", "som-ext", "ext-ext")
with_plot = False
data_path = "./DIO_6H_6_1.70bis_2.2_3.tif"

soma_low = 0.15
soma_high = 0.7126
ext_low = 0.2  # ext_low = 9.0e-4
ext_high = 0.6  # high_ext = 8.0e-3

soma_selem = mp_.disk(2)


# --- Images
image = io_.imread(data_path)
image = cl_.rgb2gray(image)[:, 512:, 512:]
img_shape = image.shape

image_for_soma = soma_.NormalizedImage(image)
image_for_ext = ext_.NormalizedImage(image)
costs = 1.0 / (image + 1.0)


# --- Initialization
n_somas = 0
som_nfo = {}  # som=soma, nfo=info
somas = None  # Tuple of soma objects

n_extensions = 0
ext_nfo = {}  # ext=extension, nfo=info
extensions = None  # Tuple of extension objects

axes = None


# --- Somas
if "soma" in run:
    print("--- Soma Detection")

    som_nfo["map"] = soma_t.Map(image_for_soma, soma_low, soma_high, soma_selem)
    som_nfo["map"] = soma_t.FilteredMap(som_nfo["map"])
    som_nfo["contour_map"] = soma_t.ContourMap(som_nfo["map"])

    som_nfo["lmp"], n_somas = ms_.label(som_nfo["map"], return_num=True)
    som_nfo["contour_lmp"] = som_nfo["contour_map"] * som_nfo["lmp"]

    som_nfo["dist_to_closest"], som_nfo["influence_map"] = soma_t.InfluenceMaps(
        som_nfo["lmp"]
    )

    somas = tuple(
        soma_t().FromMaps(som_nfo["lmp"], som_nfo["contour_lmp"], uid)
        for uid in range(1, n_somas + 1)
    )

    print(f"    n = {n_somas}")

    if with_plot:
        axes = {}
        for soma in somas:
            axes[soma.uid] = ot_.PlotLMap(som_nfo["lmp"], labels=soma.uid)
            pl_.title(f"Soma.{soma.uid}")
        # pl_.matshow(som_nfo["map"].max(axis=0)), pl_.title("Somas")
        # pl_.matshow(som_nfo["contour_map"].max(axis=0)), pl_.title("Soma Contours")
        pl_.matshow(som_nfo["influence_map"].max(axis=0)), pl_.title("Soma Influencess")
        pl_.matshow(som_nfo["dist_to_closest"].max(axis=0)), pl_.title("Soma Distances")


# -- Extentions
if "extension" in run:
    print("--- Extension Detection")

    enhanced_ext, ext_scales = extension_t.EnhancedForDetection(image_for_ext)
    ext_nfo["coarse_map"] = extension_t.CoarseMap(enhanced_ext, ext_low, ext_high)
    ext_nfo["coarse_map"] = extension_t.FilteredCoarseMap(ext_nfo["coarse_map"])
    ext_nfo["map"] = extension_t.FineMapFromCoarseMap(ext_nfo["coarse_map"])
    ext_nfo["map"][som_nfo["map"] > 0] = 0
    ext_nfo["end_point_map"] = extension_t.EndPointMap(ext_nfo["map"])

    ext_nfo["lmp"], n_extensions = ms_.label(ext_nfo["map"], return_num=True)
    ext_nfo["end_point_lmp"] = ext_nfo["end_point_map"] * ext_nfo["lmp"]

    extensions = tuple(
        extension_t().FromMaps(ext_nfo["lmp"], ext_nfo["end_point_lmp"], ext_scales, uid)
        for uid in range(1, n_extensions + 1)
    )
    for extension in extensions:
        extension.CaptureClosestSomas(som_nfo["influence_map"])

    print(f"    n = {n_extensions}")

    if with_plot:
        for extension in extensions:
            _ = ot_.PlotExtensions(extension, img_shape)
            pl_.title(f"Extension.{extension.uid}")
        # pl_.matshow(ext_nfo["map"].max(axis=0)), pl_.title("Extensions")
        pl_.matshow((10 * ext_nfo["end_point_map"] + ext_nfo["map"]).max(axis=0))
        pl_.title("Extensions Extremities")


# -- Soma-Extention
if "som-ext" in run:
    print("--- Soma <-> Extension")

    # TODO: Ideally, the extension part should be dilated
    # but in ext-ext connections, there must not be dilation around the current ext
    # (current ext plays the role of a soma in soma-ext step)
    costs[np_.logical_or(som_nfo["map"] > 0, ext_nfo["map"] > 0)] = np_.inf

    for soma in somas:
        soma.Extend(extensions, som_nfo["dist_to_closest"], costs)

    connected_ext_uids = tuple(
        extension.uid for extension in extensions if extension.soma_uid is not None
    )
    print(f"    Connected Ext = {connected_ext_uids.__len__()}"
          f"/{extensions.__len__()}\n"
          f"    {connected_ext_uids}")

    if with_plot:
        for soma in somas:
            if not soma.has_extensions:
                continue

            _ = ot_.PlotSomaWithExtensions(soma, som_nfo["lmp"], extensions)
            pl_.title(f"Soma.{soma.uid} + Ext.{soma.extension_uids}")
        soma_w_ext_lmp = soma_t.SomasWithExtensionsLMap(somas, som_nfo["lmp"], extensions)
        pl_.matshow(soma_w_ext_lmp.max(axis=0)), pl_.title("Somas + Extensions")


# -- Extention-Extention
if "ext-ext" in run:
    print("--- Extension <-> Extension")

    for extension in extensions:
        extension.Extend(extensions, costs)


elapsed_time = tm_.gmtime(tm_.time() - start_time)
print(f"\nElapsed Time={tm_.strftime('%Hh %Mm %Ss', elapsed_time)}")
print(f"DONE: {tm_.strftime('%a, %b %d %Y @ %H:%M:%S')}")

if with_plot:
    pl_.show()