#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ids_content and listidss is merged

import argparse
import sys

import numpy as np

try:
    import imaspy as imas
except ImportError:
    import imas
from rich.console import Console
from rich.markdown import Markdown
from rich.pretty import Pretty
from rich.table import Table
from rich_argparse import RichHelpFormatter

from idstools.database import DBMaster
from idstools.utils.clihelper import dbentry_parser
from idstools.utils.idshelper import (
    get_available_ids_and_occurrences,
    get_available_ids_and_times,
)
from idstools.utils.idslogger import setup_logger

logger = setup_logger("module")
np.set_printoptions(threshold=1, precision=2, suppress=True)

description = """
By default this script lists available IDSes in the pulse. It prints all IDSes along with its time length.
The fullarray option allows you to print the complete time array.
with comment option you can view ids with their occurrences, its comment and occurrence type if present
yaml option is provided which is used by other scripts to provide list of idses with their time ranges in yaml format

[``ids_content(yaml)``, ``listidss`` (with time slices), ``idsoccurrences(occ)`` merged into one script]
"""
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="idslist",
        formatter_class=RichHelpFormatter,
        description=Markdown(description, style="argparse.text"),
        parents=[dbentry_parser],
    )
    parser.add_argument(
        "-y",
        "--yaml",
        action="store_true",
        help="List ids content for yaml files aimed at describing a scenario",
    )
    parser.add_argument(
        "-c",
        "--comment",
        action="store_true",
        help="List ids with their number of occurrences with comment and occurrence_type if available",
    )
    parser.add_argument(
        "--dd-version",
        action="store_true",
        help="List ids with their number of occurrences with data dictionary version",
    )
    parser.add_argument("-i", "--ids", required=False, type=str, help="IDS name to filter")
    parser.add_argument(
        "-f",
        "--fullarray",
        action="store_true",
        help="Show full time array values [Use with default]",
    )

    args = parser.parse_args()

    connection = DBMaster.get_connection(args)
    if connection is None:
        print("----> Aborted.")
        exit(1)

    if args.yaml is True:
        if args.fullarray:
            print("(-f fullarray) option is provided, can not show full array for yaml output")
        for idsname, oc in get_available_ids_and_occurrences(connection, dd_update=args.dd_update):
            if args.ids:
                if idsname != args.ids:
                    continue

            if args.dd_update:
                ids = connection.get(idsname, occurrence=oc, autoconvert=False, ignore_unknown_dd_version=True)
                ids = imas.convert_ids(ids, connection.factory.version)
            else:
                ids = connection.get(
                    idsname, occurrence=oc, lazy=True, autoconvert=False, ignore_unknown_dd_version=True
                )
            homogeneous_time = ids.ids_properties.homogeneous_time
            if getattr(ids, "time", None) is None:
                times = None
            else:
                times = ids.time

            # Format idsname
            idsnameoc = f"{idsname}/{str(oc)}" if oc > 0 else idsname
            # For every homogeneous_time [0,1,2]
            if homogeneous_time == 0:
                print(f"  {idsnameoc}:")
                if times is not None:
                    print(f"     time_step_number: {len(times)}")
                else:
                    print(f"     time_step_number: 0")
                print("     time:             [unhomogeneous]")
            elif homogeneous_time == 1:
                if times is not None:
                    if len(times) > 1:
                        time_step = (times[len(times) - 1] - times[0]) / (len(times) - 1)
                        start_time = times[0]
                        end_time = times[len(times) - 1]
                        print(f"  {idsnameoc}:")
                        print(f"     time_step_number: {len(times)}")
                        print(f"     start_end_step:   [{str(start_time)} {str(end_time)} {str(time_step)}]")
                    elif len(times) == 1:
                        print(f"  {idsnameoc}:")
                        print(f"     time_step_number: {len(times)}")
                        print(f"     time:             [{str(times[0])}]")
            elif homogeneous_time == 2:
                print(f"  {idsnameoc}:")
                if times is not None:
                    print(f"     time_step_number: {len(times)}")
                else:
                    print(f"     time_step_number: 0")
                print("     time:             [static]")
        connection.close()
        exit(0)

    elif args.comment is True:
        console = Console()
        table = Table(title="List of IDSes with occurrences and comments")
        table.add_column("IDS", style="magenta")
        table.add_column("COMMENT", style="green")
        for idsname, oc, comment in get_available_ids_and_occurrences(
            connection, get_comment=True, dd_update=args.dd_update
        ):
            # If user wants to check specific ids type
            if args.ids:
                if idsname != args.ids:
                    continue
            # Format idsname
            if oc != 0:
                idsnameoc = f"{idsname}/{str(oc)}"
            else:
                idsnameoc = idsname
            table.add_row(idsnameoc, comment.value)
        console.print(table)
        connection.close()
        exit(0)
    elif args.dd_version is True:
        console = Console()
        table = Table(title="List of IDSes with Occurrences and DD Version")
        table.add_column("IDS", style="magenta")
        table.add_column("DD VERSION", style="green")
        for idsname, oc, dd_version in get_available_ids_and_occurrences(
            connection, get_version=True, dd_update=args.dd_update
        ):
            # If user wants to check specific ids type
            if args.ids:
                if idsname != args.ids:
                    continue
            # Format idsname
            if oc != 0:
                idsnameoc = f"{idsname}/{str(oc)}"
            else:
                idsnameoc = idsname
            table.add_row(idsnameoc, dd_version)
        console.print(table)
        connection.close()
        exit(0)
    else:
        available_ids_and_times = get_available_ids_and_times(connection, dd_update=args.dd_update)
        not_applicable = "N/A"
        console = Console()
        table = Table(title="List of IDSes")
        table.add_column("IDS")
        table.add_column("SLICES")
        table.add_column("TIME")
        for ids_name, time_array in available_ids_and_times:
            if args.ids:
                if ids_name != args.ids:
                    continue
            if time_array is not None and len(time_array) == 1 and np.isnan(time_array[0]):
                value = f"{not_applicable}"
                type = "heterogeneous IDS"
            elif time_array is not None and len(time_array) == 1 and time_array[0] == -np.inf:
                value = f"{not_applicable}"
                type = "time independent IDS"
            elif time_array is not None:
                value = f"{len(time_array)}"
                type = time_array
            elif time_array is None:
                value = f"{not_applicable}"
                type = "unknown"
            table.add_row(ids_name, value, Pretty(type))
        if args.fullarray is True:
            with np.printoptions(threshold=sys.maxsize, linewidth=1024, precision=4):
                console.print(table)
        else:
            console.print(table)
        connection.close()
        exit(0)
