import connection as cn_
import extension as xt_
import feedback as fb_
from main_prm import *
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_


soma_t = sm_.soma_t
extension_t = xt_.extension_t


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


# --- 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
som_nfo = {}  # som=soma, nfo=info
ext_nfo = {}  # ext=extension, nfo=info
axes = {}

# n_somas = 0
# n_extensions = 0
# somas = None  # Tuple of soma objects
# extensions = None  # Tuple of extension objects


# --- Somas
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
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
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.is_unconnected:
        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
print("\n--- Extension <-> Extension")

should_look_for_connections = True
while should_look_for_connections:
    som_nfo["soma_w_ext_lmp"] = soma_t.SomasLMap(somas)
    som_nfo["dist_to_closest"], som_nfo["influence_map"] = soma_t.InfluenceMaps(
        som_nfo["soma_w_ext_lmp"]
    )

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

    should_look_for_connections = False
    for ep_idx, soma, extension, end_point in candidate_conn_nfo:
        if extension.is_unconnected:
            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()