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

# Edgeplot is added in develop2.0 developed by Anna
# https://confluence.iter.org/display/IMP/How+to+plot+data+from+edge_profiles
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.compute.wall import WallCompute
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.edge_profiles import EdgeProfilesView

logger = setup_logger("module", stdout_level=logging.INFO)
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="---- Edge Profile plot",
        parents=[dbentry_parser, rcparam_parser],
        formatter_class=RichHelpFormatter,
    )

    parser.add_argument("-t", "--time", type=float, help="Time", default=-99)
    # parser.add_argument(
    #     "--separatrix",
    #     help="Show separtix",
    #     action="store_true",
    # )
    parser.add_argument(
        "--wall",
        help="Show wall",
        action="store_true",
    )
    parser.add_argument(
        "-w",
        "--walluri",
        type=str,
        help="Show wall ids from this uri (e.g. imas:hdf5?path=./testdb or ./testpulse.nc)",
    )
    parser.add_argument(
        "--logscale",
        help="Shows y axis with logarithmic scale wherever appropriate",
        action="store_true",
    )
    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()

    logger = setup_logger("module")

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

    wall_ids = None
    if args.wall:
        if args.walluri:
            mdargs = argparse.Namespace()
            mdargs.uri = args.walluri
            mdconnection = DBMaster.get_connection(args)
            try:

                if args.dd_update:
                    wall_ids = mdconnection.get("wall", autoconvert=False)
                    wall_ids = imas.convert_ids(wall_ids, mdconnection.factory.version)
                else:
                    wall_ids = mdconnection.get("wall", lazy=True, autoconvert=False)
            except Exception as e:
                logger.error(
                    f"wall ids is not present in the data entry. please provide"
                    f"with --walluri option where wall ids is exists detailed error {e}"
                )
        else:
            try:
                if args.dd_update:
                    wall_ids = connection.get("wall", autoconvert=False)
                    wall_ids = imas.convert_ids(wall_ids, connection.factory.version)
                else:
                    wall_ids = connection.get("wall", lazy=True, autoconvert=False)
            except Exception as e:
                logger.error(
                    "wall ids is not present in the data entry. please provide with"
                    f"--walluri option where wall ids is exists {e}"
                )

        if wall_ids and not wall_ids.ids_properties.homogeneous_time == imas.ids_defs.EMPTY_INT:
            r = None
            z = None
            h = None
            wall_compute = WallCompute(wall_ids)
            if vessel_units := wall_compute.get_vessel_units(name_filter="VV IS"):
                for _, description2d in vessel_units.items():
                    for v_index, vessel_unit in description2d["vesselunits"].items():
                        r = vessel_unit["r"]
                        z = vessel_unit["z"]
                        h = vessel_unit["h"]
                        break

    edge_profile_ids = None
    try:

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

    time_array = edge_profile_ids.time
    if time_array is None:
        logger.error("edge_profiles ids : time is not present")
        exit(1)

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

    if len(edge_profile_ids.grid_ggd) == 0:
        logger.error(
            f"No values in grid_ggd  at time {time_value}, Can not plot! Choose different time with option --time,"
            f"timerange [{time_array.min()}-{time_array.max()}]"
        )
        exit(1)

    if len(edge_profile_ids.grid_ggd[time_slice].space) == 0:
        logger.error(
            f"No values in grid_ggd[{time_slice}].space at time {time_value}, Can not plot!"
            f"Choose different time with option --time, timerange [{time_array.min()}-{time_array.max()}]"
        )
        exit(1)

    edge_profiles_view = EdgeProfilesView(edge_profile_ids)
    canvas = PlotCanvas(2, 4)
    canvas.update_style(args.rc)
    ax4 = canvas.add_axes(title="", xlabel="", row=0, col=0)
    ax1 = canvas.add_axes(title="", xlabel="", row=0, col=1, rowspan=2)
    ax2 = canvas.add_axes(title="", xlabel="", row=0, col=2, rowspan=2)
    ax3 = canvas.add_axes(title="", xlabel="", row=0, col=3, rowspan=2)

    c1 = edge_profiles_view.view_electrons_density(ax1, time_slice, show_separatrix=False)
    if c1 is not None:
        canvas.fig.colorbar(c1, ax=ax1)
        if wall_ids:
            if r is not None and z is not None:
                ax1.plot(r, z, "silver")

    c2 = edge_profiles_view.view_ion_density(ax2, time_slice, show_separatrix=False)
    if c2 is not None:
        canvas.fig.colorbar(c2, ax=ax2)
        if wall_ids:
            if r is not None and z is not None:
                ax2.plot(r, z, "silver")

    c3 = edge_profiles_view.view_neutral_density(ax3, time_slice, show_separatrix=False)
    if c3 is not None:
        canvas.fig.colorbar(c3, ax=ax3)
        if wall_ids:
            if r is not None and z is not None:
                ax3.plot(r, z, "silver")

    edge_profiles_view.view_equatorial_plane_and_diverter_density(ax4, time_slice, logscale=args.logscale)

    canvas.set_text(text=f"{get_database_path(args, time_value=time_value)}")

    canvas.fig.subplots_adjust(top=0.93, bottom=0.067, left=0.026, right=0.953, hspace=0.287, wspace=0.12)

    canvas.fig.suptitle(get_title(args, "Edge Profiles", time_value))
    canvas.get_current_fig_manager().set_window_title(os.path.basename(__file__))

    if args.save:
        fname = get_file_name(args, os.path.basename(__file__) + "_edge_profiles", 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()
