"""
N.B imports are within functions to prevent tensorflow being imported before
it's warnings are silenced
"""

import os
import logging
from imlib.general.logging import suppress_specific_logs

tf_suppress_log_messages = [
    "multiprocessing can interact badly with TensorFlow"
]


def main(
    signal_array,
    background_array,
    voxel_sizes,
    start_plane=0,
    end_plane=-1,
    trained_model=None,
    model_weights=None,
    model="resnet50_tv",
    batch_size=32,
    n_free_cpus=2,
    network_voxel_sizes=[5, 1, 1],
    soma_diameter=16,
    ball_xy_size=6,
    ball_z_size=15,
    ball_overlap_fraction=0.6,
    log_sigma_size=0.2,
    n_sds_above_mean_thresh=10,
    soma_spread_factor=1.4,
    max_cluster_size=100000,
    cube_width=50,
    cube_height=50,
    cube_depth=20,
    network_depth="50",
):
    suppress_tf_logging(tf_suppress_log_messages)

    from cellfinder_core.detect import detect
    from cellfinder_core.classify import classify
    from cellfinder_core.tools import prep
    from pathlib import Path

    home = Path.home()
    install_path = home / ".cellfinder"
    logging.info("Detecting cell candidates")

    points = detect.main(
        signal_array,
        start_plane,
        end_plane,
        voxel_sizes,
        soma_diameter,
        max_cluster_size,
        ball_xy_size,
        ball_z_size,
        ball_overlap_fraction,
        soma_spread_factor,
        n_free_cpus,
        log_sigma_size,
        n_sds_above_mean_thresh,
    )

    model_weights = prep.prep_classification(
        trained_model, model_weights, install_path, model, n_free_cpus
    )
    if len(points) > 0:
        logging.info("Running classification")
        points = classify.main(
            points,
            signal_array,
            background_array,
            n_free_cpus,
            voxel_sizes,
            network_voxel_sizes,
            batch_size,
            cube_height,
            cube_width,
            cube_depth,
            trained_model,
            model_weights,
            network_depth,
        )
    else:
        logging.info("No candidates, skipping classification")
    return points
    # logging.info("Saving classified cells")
    # save_cells(points, classified_points_path)


def suppress_tf_logging(tf_suppress_log_messages):
    """
    Prevents many lines of logs such as:
    "2019-10-24 16:54:41.363978: I tensorflow/stream_executor/platform/default
    /dso_loader.cc:44] Successfully opened dynamic library libcuda.so.1"
    """
    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

    for message in tf_suppress_log_messages:
        suppress_specific_logs("tensorflow", message)

[((47, 11, 47, 22), 'pathlib.Path.home', 'Path.home', ({}, {}), '()', False, 'from pathlib import Path\n'), ((49, 4, 49, 45), 'logging.info', 'logging.info', ({(49, 17, 49, 44): '"""Detecting cell candidates"""'}, {}), "('Detecting cell candidates')", False, 'import logging\n'), ((51, 13, 65, 5), 'cellfinder_core.detect.detect.main', 'detect.main', ({(52, 8, 52, 20): 'signal_array', (53, 8, 53, 19): 'start_plane', (54, 8, 54, 17): 'end_plane', (55, 8, 55, 19): 'voxel_sizes', (56, 8, 56, 21): 'soma_diameter', (57, 8, 57, 24): 'max_cluster_size', (58, 8, 58, 20): 'ball_xy_size', (59, 8, 59, 19): 'ball_z_size', (60, 8, 60, 29): 'ball_overlap_fraction', (61, 8, 61, 26): 'soma_spread_factor', (62, 8, 62, 19): 'n_free_cpus', (63, 8, 63, 22): 'log_sigma_size', (64, 8, 64, 31): 'n_sds_above_mean_thresh'}, {}), '(signal_array, start_plane, end_plane, voxel_sizes,\n    soma_diameter, max_cluster_size, ball_xy_size, ball_z_size,\n    ball_overlap_fraction, soma_spread_factor, n_free_cpus, log_sigma_size,\n    n_sds_above_mean_thresh)', False, 'from cellfinder_core.detect import detect\n'), ((67, 20, 69, 5), 'cellfinder_core.tools.prep.prep_classification', 'prep.prep_classification', ({(68, 8, 68, 21): 'trained_model', (68, 23, 68, 36): 'model_weights', (68, 38, 68, 50): 'install_path', (68, 52, 68, 57): 'model', (68, 59, 68, 70): 'n_free_cpus'}, {}), '(trained_model, model_weights, install_path, model,\n    n_free_cpus)', False, 'from cellfinder_core.tools import prep\n'), ((71, 8, 71, 46), 'logging.info', 'logging.info', ({(71, 21, 71, 45): '"""Running classification"""'}, {}), "('Running classification')", False, 'import logging\n'), ((72, 17, 86, 9), 'cellfinder_core.classify.classify.main', 'classify.main', ({(73, 12, 73, 18): 'points', (74, 12, 74, 24): 'signal_array', (75, 12, 75, 28): 'background_array', (76, 12, 76, 23): 'n_free_cpus', (77, 12, 77, 23): 'voxel_sizes', (78, 12, 78, 31): 'network_voxel_sizes', (79, 12, 79, 22): 'batch_size', (80, 12, 80, 23): 'cube_height', (81, 12, 81, 22): 'cube_width', (82, 12, 82, 22): 'cube_depth', (83, 12, 83, 25): 'trained_model', (84, 12, 84, 25): 'model_weights', (85, 12, 85, 25): 'network_depth'}, {}), '(points, signal_array, background_array, n_free_cpus,\n    voxel_sizes, network_voxel_sizes, batch_size, cube_height, cube_width,\n    cube_depth, trained_model, model_weights, network_depth)', False, 'from cellfinder_core.classify import classify\n'), ((88, 8, 88, 62), 'logging.info', 'logging.info', ({(88, 21, 88, 61): '"""No candidates, skipping classification"""'}, {}), "('No candidates, skipping classification')", False, 'import logging\n'), ((103, 8, 103, 53), 'imlib.general.logging.suppress_specific_logs', 'suppress_specific_logs', ({(103, 31, 103, 43): '"""tensorflow"""', (103, 45, 103, 52): 'message'}, {}), "('tensorflow', message)", False, 'from imlib.general.logging import suppress_specific_logs\n')]