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

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,
    get_database_path,
    get_file_name,
    get_title,
    rcparam_parser,
)
from idstools.utils.idslogger import setup_logger
from idstools.view.common import PlotCanvas
from idstools.view.core_sources import CoreSourcesView

logger = setup_logger("module", stdout_level=logging.INFO)
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="---- Display core_sources results [Previously known as csplot]",
        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 figure at default location",
        action="store_true",
    )
    parser.add_argument(
        "--directory",
        help="Directory to save the figure",
        default=None,
    )

    args = parser.parse_args()
    connection = DBMaster.get_connection(args)
    if connection is None:
        logger.critical("----> Aborted.")
        exit(1)

    core_sources = None

    try:

        if args.dd_update:
            core_sources = connection.get("core_sources", autoconvert=False)
            core_sources = imas.convert_ids(core_sources, connection.factory.version)
        else:
            core_sources = connection.get("core_sources", lazy=True, autoconvert=False)

    except Exception as e:
        logger.error(f"core_sources ids is not present {e}")
        exit(1)

    if core_sources.time is None:
        logger.critical("No core_sources IDS in the data-entry --> Abort.")
        exit()

    time_array = core_sources.time
    time_slice, time_value = get_nearest_time(time_array, args.time)
    ntime = len(time_array)
    nsources = len(core_sources.source)

    core_source_view = CoreSourcesView(core_sources)

    canvas = PlotCanvas(2, 4)
    canvas.update_style(args.rc)
    ret = ax_power_profiles = canvas.add_axes(title=r"Power Profiles [MW/M3]", row=0, col=0)
    core_source_view.view_power_profiles(ax_power_profiles, time_slice)

    ax_particles_profiles = canvas.add_axes(title=r"PARTICLES PROFILES [/M3/S]", row=0, col=1)
    ret = core_source_view.view_particles_profiles(ax_particles_profiles, time_slice)

    ax_current_profiles = canvas.add_axes(title=r"CURRENT PROFILES [KA/M2]", row=0, col=2)
    ret = core_source_view.view_current_profiles(ax_current_profiles, time_slice)

    ax_power_and_particle_waveforms = canvas.add_axes(title="POWER AND PARTICLE WAVEFORMS", row=1, col=0)
    ret = core_source_view.view_power_and_particle_waveforms(ax_power_and_particle_waveforms, time_slice)

    ax_particles_waveform = canvas.add_axes(title="PARTICLES WAVEFORM", row=1, col=1)
    ret = core_source_view.view_particles_waveform(ax_particles_waveform, time_slice)

    ax_current_waveform = canvas.add_axes(title=r"CURRENT WAVEFORM", row=1, col=2)
    ret = core_source_view.view_current_waveform(ax_current_waveform, time_slice)

    ax_torque_waveform = canvas.add_axes(row=1, col=3)
    ret = core_source_view.view_torque_waveform(ax_torque_waveform, time_slice)

    canvas.set_text(text=f"{get_database_path(args, time_value=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.fig.suptitle(get_title(args, "Core sources", time_value))
    if args.save:
        fname = get_file_name(args, os.path.basename(__file__) + "_core_sources", time_value)
        if args.directory:
            if not os.path.exists(args.directory):
                os.makedirs(args.directory)
            fname = os.path.join(args.directory, fname)
        canvas.save(fname)
    else:
        canvas.show()
