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

# $ ids_compo -p 131047 -r 4
# No edge_profiles IDS in the data-entry.
# core +  edge  -
#    ------------
# core_profiles
#    ------------
#    species:      H         D         T         He3       He4       Be        Ne
#    a:            1.0       2.0       3.0       3.0       4.0       9.0       20.0
#    z:            1.0       1.0       1.0       2.0       2.0       4.0       10.0
#    n_over_ntot:  5.29e-06  0.460     0.493     7.01e-07  0.011     0.024     0.012
#    n_over_ne:    4.45e-06  0.387     0.414     5.89e-07  9.58e-03  0.020     0.010
#    n_over_n_maj: 1.07e-05  0.933     1.000     1.42e-06  0.023     0.048     0.024
import argparse
import logging

import numpy as np

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
from idstools.utils.idslogger import setup_logger
from idstools.view.core_profiles import CoreProfilesView
from idstools.view.edge_profiles import EdgeProfilesView

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="---- Display the plasma composition from the core_profiles IDS [previously known as ids_compo]",
        parents=[dbentry_parser],
        formatter_class=RichHelpFormatter,
    )
    parser.add_argument("-t", "--time", help="Time", required=False, type=float, default=-99.0)
    parser.add_argument("-i", "--info", action="store_true", help="Show information")
    parser.add_argument("--debug", action="store_true", help="Show debugging")

    args = parser.parse_args()

    level = logging.ERROR
    if args.info:
        level = logging.INFO
    if args.debug:
        level = logging.DEBUG

    logger = setup_logger("module", stdout_level=level)

    if args.info:
        logger.info("----------------------------------------------------------------")
        logger.info(
            """
            a       = Mass of atom [Atomic Mass Unit]
            z       = Nuclear charge [Elementary Charge Unit]
            ne      = sum(volume * electron_density)
            volume  = sum(volume of each cell)
            nspec_over_ntot : Species density list / Sum of Species Density (ntot)
            nspec_over_ne   : Species density list / Total no. electrons (ne)
            nspec_over_nmaj : Species density list / Species density[Index of Maximum Density Species]"""
        )
        logger.info("----------------------------------------------------------------")
    connection = DBMaster.get_connection(args)
    if connection is None:
        logger.critical("----> Aborted.")
        exit(1)

    # Prepare IDS with highest density data from core profiles
    try:
        logger.info("\ngetting core_profiles slice")

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

        if args.time == -99:
            ne0_list = np.array([])
            for i, _ in enumerate(core_profiles_ids.time):
                ne0_list = np.append(ne0_list, core_profiles_ids.profiles_1d[i].electrons.density[0])
            time_slice_core_profiles = np.argmax(ne0_list)
            time_value_core_profiles = core_profiles_ids.time[time_slice_core_profiles]
            logger.info(
                f"core_profiles IDS:Using time slice where maximum electrons density is present,"
                f"Time slice:{time_slice_core_profiles}"
            )
        else:
            time_array_core_profiles = core_profiles_ids.time
            time_slice_core_profiles, time_value_core_profiles = get_nearest_time(time_array_core_profiles, args.time)
        logger.info(f"coreprofiles Time {time_value_core_profiles}  Index {time_slice_core_profiles}")

    except Exception as e:
        logger.error(f"Exception occurred, detailed error {e}")
        core_profiles_ids = None

    try:
        logger.info("\ngetting edge_profiles slice")

        if args.dd_update:
            edge_profiles_ids = connection.get("edge_profiles", autoconvert=False)
            edge_profiles_ids = imas.convert_ids(edge_profiles_ids, connection.factory.version)
        else:
            edge_profiles_ids = connection.get("edge_profiles", lazy=True, autoconvert=False)
        if time_slice_core_profiles is None:
            if args.time == -99:
                ne0_list = np.array([])
                for i, _ in enumerate(edge_profiles_ids.time):
                    ne0_list = np.append(ne0_list, edge_profiles_ids.profiles_1d[i].electrons.density[0])
                time_slice_edge_profiles = np.argmax(ne0_list)
                time_value_edge_profiles = edge_profiles_ids.time[time_slice_edge_profiles]
                logger.info(
                    f"edge_profiles IDS:Using time slice where maximum electrons density is present,"
                    f"Time slice:{time_slice_edge_profiles}"
                )
            else:
                time_array_edge_profiles = edge_profiles_ids.time
                time_slice_edge_profiles, time_value_edge_profiles = get_nearest_time(
                    time_array_edge_profiles, args.time
                )
        else:
            time_array_edge_profiles = edge_profiles_ids.time
            time_slice_edge_profiles, time_value_edge_profiles = get_nearest_time(
                time_array_edge_profiles, time_value_core_profiles
            )
        logger.info(f"edge_profiles Time {time_value_edge_profiles}  Index {time_slice_edge_profiles}")

    except Exception as e:
        logger.error(f"Exception occurred, detailed error {e}")
        edge_profiles_ids = None

    core_profile_exists = True
    edge_profile_exists = True
    if core_profiles_ids is None:
        core_profile_exists = False
    if edge_profiles_ids is None:
        edge_profile_exists = False

    if core_profiles_ids is not None:

        returnstatus = CoreProfilesView.view_plasma_composition_with_species_concentration(
            core_profiles_ids, time_slice_core_profiles
        )
        if returnstatus == 0:
            core_profile_exists = False
            logger.warning("core_profiles: IDS exists but time slice doesn't exists.")
        elif returnstatus == -1:
            core_profile_exists = False
            logger.warning(
                f"core_profiles: IDS exists but volume is not set, " f"try retrieving it from the equilibrium IDS"
            )
            equilibrium_ids = None
            try:

                if args.dd_update:
                    equilibrium_ids = connection.get("equilibrium", autoconvert=False)
                    equilibrium_ids = imas.convert_ids(equilibrium_ids, connection.factory.version)
                else:
                    equilibrium_ids = connection.get("equilibrium", lazy=True, autoconvert=False)
                volume = equilibrium_ids.time_slice({time_slice_core_profiles}).profiles_1d.volume
            except Exception as e:
                logger.warning(f"equilibrium: IDS exists but volume is not present. detailed error {e}")

            returnstatus = CoreProfilesView.view_plasma_composition_with_species_concentration(
                core_profiles_ids, time_slice_core_profiles, volume=volume
            )
            core_profile_exists = True
    else:
        logger.warning("No core_profiles IDS in the data-entry.")

    if edge_profiles_ids is not None:
        logger.info(
            f"edge_profiles: using time slice where maximum "
            f"electrons density is present in core, time slice={time_slice_edge_profiles}"
        )

        print("----------------")
        print("edge_profiles")
        print("----------------")
        returnstatus = EdgeProfilesView.view_plasma_composition_with_species_concentration(
            edge_profiles_ids, time_slice_edge_profiles
        )
        if returnstatus == 0:
            edge_profile_exists = False
            logger.warning("edge_profiles: time slice doesn't exists.")
        elif returnstatus == -1:
            logger.warning(
                f"edge_profiles: IDS exists but volume is not set. " f"Trying nearby time slice. This may take a while."
            )
            edge_profiles = connection.get("edge_profiles", autoconvert=False)
            if args.dd_update:
                edge_profiles = imas.convert_ids(edge_profiles, connection.factory.version)
            time_slices_count = len(edge_profiles.grid_ggd)
            index_list = []
            for counter in range(1, 10):
                if time_slice_edge_profiles - counter >= 0 and time_slice_edge_profiles - counter < time_slices_count:
                    index_list.append(time_slice_edge_profiles - counter)
                if time_slice_edge_profiles + counter >= 0 and time_slice_edge_profiles + counter < time_slices_count:
                    index_list.append(time_slice_edge_profiles + counter)
            for itime_slice in index_list:
                index = 4
                elements = edge_profiles.grid_ggd[itime_slice].grid_subset[index].element
                grid_subset_name = edge_profiles.grid_ggd[itime_slice].grid_subset[index].identifier.name
                # check if grid_subset[4] identifier name is cells, if not, find out 'cells' index
                index_counter = 0
                if grid_subset_name.lower() != "cells":
                    for subset in edge_profiles.grid_ggd[itime_slice].grid_subset:
                        if subset.identifier.name.lower() == "cells":
                            elements = edge_profiles.grid_ggd[itime_slice].grid_subset[index_counter].element
                            grid_subset_name = (
                                edge_profiles.grid_ggd[itime_slice].grid_subset[index_counter].identifier.name
                            )
                            index = index_counter
                            break
                        index_counter = index_counter + 1
                if len(elements) != 0:
                    logger.info(
                        f"edge_profiles: using nearby time slice in edge where maximum "
                        f"electrons density is present in core, Time slice:{itime_slice}"
                    )
                    returnstatus = EdgeProfilesView.view_plasma_composition_with_species_concentration(
                        edge_profiles_ids, itime_slice
                    )
                    if returnstatus == 0:
                        edge_profile_exists = False
                        logger.warning("edge_profiles: time slice doesn't exists.")
                    elif returnstatus == -1:
                        edge_profile_exists = False
                        logger.warning("edge_profiles: IDS exists but volume is not set.")
                    break
    else:
        logger.warning("No edge_profiles IDS in the data-entry.")

    profile_availability_string = "" + ("core -\t" if not core_profile_exists else "core +\t")
    if edge_profile_exists is False:
        profile_availability_string += "edge -\t"
    else:
        profile_availability_string += "edge +\t"

    print(profile_availability_string)
