#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# migrated from neutronplot
import argparse
import logging
import os

try:
    import imaspy as imas

except ImportError:
    import imas
from rich_argparse import RichHelpFormatter

from idstools.compute.common import get_nearest_time
from idstools.database import DBMaster
from idstools.utils.clihelper import (
    dbentry_parser,
    show_plot,
    get_database_path,
    get_file_name,
    rcparam_parser,
)
from idstools.utils.idslogger import setup_logger
from idstools.view.common import PlotCanvas
from idstools.view.distribution_sources import DistributionSourcesView

if __name__ == "__main__":
    # Management of input arguments
    parser = argparse.ArgumentParser(
        description="""---- Display the neutron profiles from the distribution_sources IDSs
        [previously known as neutronplot]""",
        parents=[dbentry_parser, rcparam_parser],
        formatter_class=RichHelpFormatter,
    )

    parser.add_argument("-t", "--time", help="Time", required=False, type=float, default=-99.0)
    parser.add_argument(
        "--save",
        help=(
            "Save the plot to a file instead of opening a window. "
            "If no plot window is available, it is saved automatically."
        ),
        action="store_true",
    )
    parser.add_argument(
        "--directory",
        help="Directory to save the figure",
        default=None,
    )
    args = parser.parse_args()
    logger = setup_logger("module", stdout_level=logging.INFO)

    # Open the database and read the necessary IDSs
    connection = DBMaster.get_connection(args)
    if connection is None:
        logger.critical("----> Aborted.")
        exit(1)

    hostdir = get_database_path(args)
    if hostdir is None:
        logger.critical(
            "Environment variable IMAS_HOME is not defined. Quitting.",
        )
        exit(1)

    ids_distribution_sources = None
    try:

        if args.dd_update:
            ids_distribution_sources = connection.get("distribution_sources", autoconvert=False)
            ids_distribution_sources = imas.convert_ids(ids_distribution_sources, connection.factory.version)
        else:
            ids_distribution_sources = connection.get("distribution_sources", lazy=True, autoconvert=False)
    except Exception as e:
        logger.error(f"distribution_sources ids is not present, detailed error: {e}")
        exit(1)

    time_array = ids_distribution_sources.time
    time_slice, time_value = get_nearest_time(time_array, args.time)

    nsources = len(ids_distribution_sources.source)
    if nsources > 1:
        logger.info(
            f"Distribution_sources contains {nsources} sources",
        )
    else:
        logger.info(
            f"Distribution_sources contains {nsources} source",
        )

    active = [1] * nsources  # A priori all sources are active
    for isource in range(nsources):
        try:
            availability_test = ids_distribution_sources.source[isource].profiles_1d[time_slice]

        except Exception as e:
            logger.warning(
                f"Source {str(isource)} has no profiles_1d -> considered as inactive, detailed error: {e}",
            )
            active[isource] = 0
        try:
            availability_test = ids_distribution_sources.source[isource].global_quantities[time_slice]
        except Exception as e:
            if active[isource] == 1:
                logger.warning(
                    f"Source {str(isource)} has no global_quantities -> considered as inactive, detailed error: {e}",
                )
                active[isource] = 0
        try:
            availability_test = ids_distribution_sources.source[isource].process[time_slice]
        except Exception as e:
            if active[isource] == 1:
                logger.warning(
                    f"Source {str(isource)} has no process -> considered as inactive, detailed error: {e}",
                )
                active[isource] = 0

    if sum(active) == 0:
        logger.critical("No active neutron sources --> Leave.")
        exit()

    canvas = PlotCanvas(1, 1)
    canvas.update_style(args.rc)
    ax = canvas.add_axes(title="", xlabel="", row=0, col=0)
    distribution_sources_view = DistributionSourcesView(ids_distribution_sources)
    distribution_sources_view.view_neutrons(ax, time_slice)
    distribution_sources_view.view_time(ax, time_value)

    canvas.fig.subplots_adjust(top=0.916, bottom=0.09, left=0.044, right=0.953, hspace=0.287, wspace=0.2)
    canvas.get_current_fig_manager().set_window_title(os.path.basename(__file__))
    canvas.set_sup_title(get_database_path(args, time_value=time_value))

    show_plot(canvas, args, os.path.basename(__file__) + "_Neutrons", time_value)
