import connection as cn_
import extension as xt_
import feedback as fb_
import soma as sm_

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 = sm_.soma_t
extension_t = xt_.extension_t


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

soma_low_c = 0.15
soma_high_c = 0.7126
ext_low_c = 0.2  # 0.02  # 0.2  # ext_low_c = 9.0e-4
ext_high_c = 0.6  # 0.04  # 0.6  # high_ext = 8.0e-3

soma_selem_c = mp_.disk(2)

max_straight_sq_dist_c = 30 ** 2
max_weighted_length_c = 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 = sm_.NormalizedImage(image)
image_for_ext = xt_.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 = {}


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

    som_nfo["map"] = soma_t.Map(image_for_soma, soma_low_c, soma_high_c, soma_selem_c)
    som_nfo["map"] = soma_t.FilteredMap(som_nfo["map"], min_area_c)
    som_nfo["lmp"], n_somas = ms_.label(som_nfo["map"], return_num=True)
    som_nfo["dist_to_closest"], som_nfo["influence_map"] = soma_t.InfluenceMaps(
        som_nfo["lmp"]
    )

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

    print(f"    n = {n_somas}")
    if with_plot:
        fb_.PlotSomas(somas, som_nfo, axes)


# -- 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_c, ext_high_c)
    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["lmp"], n_extensions = ms_.label(ext_nfo["map"], return_num=True)

    extensions = tuple(
        extension_t().FromMap(ext_nfo["lmp"], ext_scales, uid)
        for uid in range(1, n_extensions + 1)
    )

    print(f"    n = {n_extensions}")
    if with_plot:
        fb_.PlotExtensions(extensions, ext_nfo, img_shape)


# -- 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 = cn_.CandidateConnections(
        somas,
        som_nfo["influence_map"],
        som_nfo["dist_to_closest"],
        extensions,
        max_straight_sq_dist_c,
    )

    for ep_idx, soma, extension, end_point in candidate_conn_nfo:
        if extension.soma_uid is not None:
            continue

        print(f"    Soma.{soma.uid} <-?-> Ext.{extension.uid}({ep_idx})", end="")
        path, length = cn_.ShortestPathTo(
            end_point,
            costs,
            soma.ContourPointsCloseTo,
            max_straight_sq_dist=max_straight_sq_dist_c,
        )
        if length <= max_weighted_length_c:
            soma.ExtendWith(extension, path, costs)
            print(": Made")
        else:
            print("")

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

    fb_.PrintConnectedExtensions(extensions)
    if with_plot:
        fb_.PlotSomasWithExtensions(somas, som_nfo, "all")


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

        candidate_conn_nfo = cn_.CandidateConnections(
            somas,
            som_nfo["influence_map"],
            som_nfo["dist_to_closest"],
            unconnected_extensions,
            max_straight_sq_dist_c,
        )

        should_look_for_connections = False
        for ep_idx, soma, extension, end_point in candidate_conn_nfo:
            if extension.soma_uid is not None:
                continue

            print(f"    Soma.{soma.uid} <-?-> Ext.{extension.uid}({ep_idx})", end="")
            path, length = cn_.ShortestPathTo(
                end_point,
                costs,
                soma.ExtensionPointsCloseTo,
                max_straight_sq_dist=max_straight_sq_dist_c,
            )
            if length <= max_weighted_length_c:
                tgt_extenstion = extension_t.ExtensionWithSite(extensions, path[-1])
                tgt_extenstion.ExtendWith(extension, path, costs)
                should_look_for_connections = True
                print(": Made")
            else:
                print("")

    fb_.PrintConnectedExtensions(extensions)
    if with_plot:
        fb_.PlotSomasWithExtensions(somas, som_nfo, "with_ext_of_ext")


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