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

import itertools as it_
import time as tm_

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_


print(f"STARTED: {tm_.strftime('%a, %b %d %Y @ %H:%M:%S')}")
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)

max_straight_sq_dist = 30 ** 2
max_weighted_length = 20.0

min_area_c = 1000


# --- 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("\n--- 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"], min_area_c)
    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("\n--- 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("\n--- 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

    candidate_conn_nfo = []  # conn=connection
    for soma, extension in it_.product(somas, extensions):
        new_candidates = extension.EndPointsForSoma(soma.uid)
        candidate_conn_nfo.extend(
            (soma, extension, end_point) for end_point in new_candidates
        )
    candidate_conn_nfo.sort(key=lambda elm: som_nfo["dist_to_closest"][elm[2]])

    for c_idx, (soma, extension, end_point) in enumerate(candidate_conn_nfo):
        print(f"    Soma.{soma.uid} <-?-> Ext.{extension.uid}({c_idx})", end="")
        path, length = soma_t.ShortestPathFrom(
            end_point,
            costs,
            soma.ContourPointsCloseTo,
            max_straight_sq_dist=max_straight_sq_dist,
        )
        if length <= max_weighted_length:
            soma.ExtendWith(extension, path, costs)
            print(": Made")
        else:
            print("")
    som_nfo["soma_w_ext_lmp"] = soma_t.SomasWithExtensionsLMap(somas)

    # 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:
        # Somas without extensions have already been plotted
        for soma in filter(lambda elm: elm.has_extensions, somas):
            _ = ot_.PlotSomaWithExtensions(soma, som_nfo["lmp"])
            pl_.title(
                f"Soma.{soma.uid} + Ext.{tuple(ext.uid for ext in soma.extensions)}"
            )
        pl_.matshow(som_nfo["soma_w_ext_lmp"].max(axis=0)), pl_.title(
            "Somas + Extensions"
        )


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

    should_look_for_connections = True
    while should_look_for_connections:
        som_nfo["soma_w_ext_lmp"] = soma_t.SomasWithExtensionsLMap(somas)
        som_nfo["dist_to_closest"], som_nfo["influence_map"] = soma_t.InfluenceMaps(
            som_nfo["soma_w_ext_lmp"]
        )
        unconnected_extensions = list(
            filter(lambda ext: ext.soma_uid is None, extensions)
        )
        for extension in unconnected_extensions:
            extension.CaptureClosestSomas(som_nfo["influence_map"])

        candidate_conn_nfo = []  # conn=connection
        for soma, extension in it_.product(somas, unconnected_extensions):
            new_candidates = extension.EndPointsForSoma(soma.uid)
            candidate_conn_nfo.extend(
                (soma, extension, end_point) for end_point in new_candidates
            )
        candidate_conn_nfo.sort(key=lambda elm: som_nfo["dist_to_closest"][elm[2]])

        should_look_for_connections = False
        for c_idx, (soma, extension, end_point) in enumerate(candidate_conn_nfo):
            print(f"    Soma.{soma.uid} <-?-> Ext.{extension.uid}({c_idx})", end="")
            path, length = soma_t.ShortestPathFrom(
                end_point,
                costs,
                soma.ExtensionPointsCloseTo,
                max_straight_sq_dist=max_straight_sq_dist,
            )
            if length <= max_weighted_length:
                tgt_extenstion = extension_t.ExtensionWithSite(extensions, path[-1])
                tgt_extenstion.ExtendWith(extension, path, costs)
                should_look_for_connections = True
                print(": Made")
            else:
                print("")


# -- Summary
print("\n")
for soma in somas:
    print(soma)
for extension in extensions:
    print(extension)

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()