# Copyright CNRS/Inria/UNS
# Contributor(s): Eric Debreuve (since 2018)
#
# eric.debreuve@cnrs.fr
#
# This software is governed by the CeCILL  license under French law and
# abiding by the rules of distribution of free software.  You can  use,
# modify and/ or redistribute the software under the terms of the CeCILL
# license as circulated by CEA, CNRS and INRIA at the following URL
# "http://www.cecill.info".
#
# As a counterpart to the access to the source code and  rights to copy,
# modify and redistribute granted by the license, users are provided only
# with a limited warranty  and the software's author,  the holder of the
# economic rights,  and the successive licensors  have only  limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using,  modifying and/or developing or reproducing the
# software by the user in light of its specific status of free software,
# that may mean  that it is complicated to manipulate,  and  that  also
# therefore means  that it is reserved for developers  and  experienced
# professionals having in-depth computer knowledge. Users are therefore
# encouraged to load and test the software's suitability as regards their
# requirements in conditions enabling the security of their systems and/or
# data to be ensured and,  more generally, to use and operate it in the
# same conditions as regards security.
#
# The fact that you are presently reading this means that you have had
# knowledge of the CeCILL license and that you accept its terms.

import matplotlib as mpl_

mpl_.use('TkAgg')

from sklgraph.skl_fgraph import skl_graph_t
from sklgraph.skl_graph import plot_mode_e
from sklgraph.skl_map import skl_map_t

import pathlib as ph_
import sys as sy_
import tempfile as tf_

import matplotlib.pyplot as pl_
import numpy as np_
import scipy.ndimage as im_
import skimage.measure as ms_

# import sknw as sk_


# --- Set main parameters
if sy_.argv.__len__() == 1:
    print("Call with arguments: horse|bug|2d|3d [noplot]")
    sy_.exit(0)

data = sy_.argv[1]
print(data)
test_mode = (sy_.argv.__len__() > 2) and (sy_.argv[2] == "noplot")
space_dim = None

# --- Generation of a (random) object map
if data == "horse":
    import skimage.data as sd_
    import skimage.util as su_

    obj_map = su_.invert(sd_.horse())
    space_dim = 2
#
elif data == "bug":
    bug_path = ph_.Path(__file__).parent
    bug_path = bug_path / ".." / "dev" / "__usecases__" / "obj_map.npz"
    if bug_path.is_file():
        obj_map = np_.load(bug_path)["obj_map"]
        space_dim = obj_map.ndim
    else:
        raise FileNotFoundError(bug_path)
#
else:
    if data == "2d":
        space_dim = 2
        array_shape = (100, 100)
        threshold_fct = lambda img: img.median()
        struct_elm_shape = (3, 3)
    elif data == "3d":
        space_dim = 3
        array_shape = (80, 80, 50)
        threshold_fct = lambda img: 0.75 * image.max()
        struct_elm_shape = (3, 3, 3)
    else:
        raise ValueError(f"{data}: Unknown data specification")

    tmp_folder = None
    with tf_.NamedTemporaryFile() as doc_accessor:
        doc_name = ph_.Path(doc_accessor.name)
        tmp_folder = doc_name.parents[0]
    if tmp_folder is not None:
        seed_lock = tmp_folder.joinpath("skl_graph_demo_seed.lock")
        if not seed_lock.exists():
            np_.random.seed(0)
            seed_lock.touch()

    image = np_.random.random(array_shape)
    image = im_.gaussian_filter(image, 1)

    multi_map = image > threshold_fct(image)
    lbl_multi_map, _ = im_.label(
        multi_map,
        structure=np_.ones(struct_elm_shape, dtype=np_.uint8),
        output=np_.uint,
    )

    cc_props = ms_.regionprops(lbl_multi_map)
    largest_cc_idx = np_.argmax([cc_prop.area for cc_prop in cc_props]).item()
    lbl_multi_map[lbl_multi_map != cc_props[largest_cc_idx].label] = 0
    obj_map = lbl_multi_map > 0

# --- From object map to skeleton graph
np_.savez_compressed(".\\obj_map", obj_map=obj_map)
skeleton = skl_map_t.FromShapeMap(obj_map, store_widths=True)
skel_graph = skl_graph_t.FromSkeleton(skeleton)

# sknw_graph = sk_.build_sknw(skeleton.map)
# print(dir(sknw_graph))

# --- Some info about the skeleton graph
print(
    f"Obj map area={np_.count_nonzero(obj_map)}\n\n"
    f"Validity={skel_graph.is_valid}\n\n"
    f"N nodes={skel_graph.n_nodes}\n"
    f"N edges={skel_graph.n_edges}\n"
    f"Highest degree={skel_graph.highest_degree}/{skel_graph.highest_degree_w_nodes}\n\n"
    f"Length={skel_graph.length}<-{skel_graph.edge_lengths}\n"
    f"Width=Hom.{skel_graph.reduced_width()}/Het.{skel_graph.heterogeneous_reduced_width()}<-{skel_graph.edge_reduced_widths()}\n"
    f"Area as WxL={skel_graph.reduced_width() * skel_graph.length}\n"
    f"Area as WW Length={skel_graph.ww_length}<-{skel_graph.edge_ww_lengths}\n"
)

# --- No display is testing
if test_mode:
    sy_.exit(0)

# --- Checking correctness of skeleton graph via various rebuilt maps
rebuilt_skel = skel_graph.RebuiltSkeletonMap()
part_map_for_check = skeleton.PartMap()
# Keep next line before its next one
part_map_for_check[part_map_for_check == skeleton.invalid_n_neighbors] = 0
part_map_for_check[part_map_for_check > 3] = 3

print(
    f"Rebuilt skeleton = Original one? Part map? "
    f"{np_.array_equal(rebuilt_skel > 0, skeleton.map)} "
    f"{np_.array_equal(rebuilt_skel, part_map_for_check)}"
)

if space_dim == 3:
    rebuilt_skel = rebuilt_skel.max(axis=2)

# --- Checking correctness of skeleton graph via rebuilt object map
obj_map = obj_map.astype(np_.uint8)
rebuilt_map = skel_graph.RebuiltObjectMap()
if space_dim == 3:
    obj_map = obj_map.max(axis=2)
    rebuilt_map = rebuilt_map.max(axis=2)

# --- Show various images and plot the skeleton graph
___, axes = pl_.subplots(1, 3)
axes[0].matshow(obj_map)
axes[1].matshow(rebuilt_map)
axes[2].matshow(obj_map - rebuilt_map)
axes[0].set_title("Object")
axes[1].set_title("Rebuilt")
axes[2].set_title("Object - Rebuilt")
for axis in axes:
    axis.set_axis_off()
pl_.tight_layout(pad=0.3)

___, axes = pl_.subplots(1, 3)
axes[0].matshow(rebuilt_skel + obj_map)
axes[1].matshow(rebuilt_skel + rebuilt_map)
axes[2].matshow(rebuilt_skel)
axes[0].set_title("Rebuilt Skel Over Obj")
axes[1].set_title("Rebuilt Skel Over Rebuilt Obj")
axes[2].set_title("Rebuilt Skel")
for axis in axes:
    axis.set_axis_off()
pl_.tight_layout(pad=0.3)

if space_dim == 2:
    skel_graph.Plot(mode=plot_mode_e.Networkx, should_block=False)
    skel_graph.Plot(mode=plot_mode_e.Graphviz, should_block=False)
skel_graph.Plot(mode=plot_mode_e.SKL_Curve, w_directions=True, should_block=False)
skel_graph.Plot(should_block=False)

# From: https://github.com/Image-Py/sknw (previously: https://github.com/yxdragon/sknw)
# sk_.draw_graph(obj_map, sknw_graph, cn=255, ce=128)
# pl_.matshow(obj_map, cmap="gray")
# pl_.title("SKNW")
# pl_.gca().set_axis_off()
# pl_.tight_layout(pad=0.3)

pl_.show()