#!/usr/bin/env python3
# THE ROSETTA CODE
import argparse
import math
import os
import sys

try:
    import imaspy as imas
except ImportError:
    import imas
import numpy as np
import pandas as pd
import requests
from packaging.version import Version
from rich_argparse import RichHelpFormatter

from idstools.utils.clihelper import get_backend_id
from idstools.utils.idshelper import get_ids_types


def ids_setter(ids, path, val):
    """Set a value in a field of the given IDS.

    Parameters
    ----------
    IDS  : IDS object
    path : str
           Path from root of the IDS to field in which the value shall be stored.
           Strutures are separated by slashes and array of structure indices are given between parenthesis.
           eg. 'structA/arraystruct(0)/structB/value'
    val
           Value to be stored in the IDS/path node
    """
    dodi = ids
    for node in path[0:-1]:
        nodestruct = node.split("(")

        try:
            dodi = getattr(dodi, nodestruct[0])
        except AttributeError:
            raise AttributeError(
                f"Node {node} could not be found in {path} for IDS {ids._path}. Please check spelling."
            )
        if len(nodestruct) == 2:
            # AOS case
            nodeindex = int(nodestruct[1][:1])
            dodi.resize(nodeindex + 1, keep=True)
            dodi = dodi[nodeindex]
    leaf = path[-1]

    try:
        ids_field = getattr(dodi, leaf)
        leaftype = type(ids_field.value)

        if leaftype == str:
            setattr(dodi, str(leaf), str(val))
        elif leaftype == float:
            setattr(dodi, str(leaf), float(val))
        elif leaftype == int:
            setattr(dodi, str(leaf), int(val))
        elif leaftype == np.ndarray:
            if len(ids_field.shape) == 1:
                arr = [val]
            elif len(ids_field.shape) == 2:
                arr = [[val]]
            setattr(dodi, str(leaf), np.array(arr, dtype=ids_field.dtype))
        else:
            raise AttributeError(f"The type {leaftype} of {idspath} is currently not supported")
    except AttributeError:
        raise AttributeError(
            f"The leaf '{node}' could not be found in {path} for IDS {ids._base_path}. Please check spelling."
        )


def evaluate(expr):
    """Evaluate a simple expression

    Parameter
    ---------
    expr : str
           Expression to be evaluated

    Returns
    -------
    int or float or str (None if evaluation failed)
    """
    result = None
    try:
        result = eval(expr)
    except KeyError as ke:
        print(f"Can't map key {ke} for expression '{expr}'")
    except SyntaxError as se:
        print(f"Syntax error {se} for expression '{expr}'")
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="This script applies mapping of a non-IDS database content (e.g. ITPA DBs)"
        "into IDS rules. replaced by ids_rosetta_code.py",
        formatter_class=RichHelpFormatter,
    )
    parser.add_argument(
        "-i",
        "--inputcsv",
        type=str,
        default="HDB5.2.3.csv",
        help="Path to csv file containing the external database content \t(default=%(default)s, "
        "the H-mode DB will be downloaded from https://osf.io/zhwa3/ automatically if not present)",
    )
    parser.add_argument(
        "-m",
        "--mapping",
        type=str,
        default=os.path.join(
            os.path.dirname(os.path.dirname(os.path.realpath(__file__))), "bin/mappings/h-mode-db-mapping.csv"
        ),
        help="Path to csv-formatted mapping file \t(default=%(default)s)",
    )
    parser.add_argument(
        "--varcol",
        type=str,
        default="DB_VARIABLE",
        help="Name of the column of the mapping file listing all DB variables \t(default=%(default)s)",
    )
    parser.add_argument(
        "--pathcol",
        type=str,
        default="IDS_PATH",
        help="Name of the column of the mapping file listing the paths to store all"
        "DB variables into IDS fields \t(default=%(default)s)",
    )
    parser.add_argument(
        "--tracol",
        type=str,
        default="TRANSFORMATION",
        help="Name of the column of the mapping file listing the transformations"
        "to be done on DB variables (default = summary)\t(default=%(default)s)",
    )
    parser.add_argument(
        "--timeloc",
        type=str,
        default="summary",
        help="Name of the IDS from which the time will be extracted to populate time-empty IDSs \t(default=%(default)s",
    )
    parser.add_argument(
        "--mapping-version",
        type=str,
        default="3.37.0",
        help="Data dictionary version for which the mapping was expressed",
    )
    parser.add_argument(
        "--convert-to",
        type=str,
        default=None,
        help=(
            "If enabled, convert IDS to the specified version of the data dictionary before storing."
            "Otherwise, store using the mapping-version."
        ),
    )
    parser.add_argument("-b", "--backend", type=str, default="ASCII", help="backend format \t(default=%(default)s)")
    parser.add_argument(
        "-d", "--database", type=str, default="test", help="target IMAS database name \t(default=%(default)s)"
    )
    parser.add_argument(
        "--dbtype", type=str, default="H-MODE", help="Type of database (default = H-MODE) \t(default=%(default)s)"
    )
    parser.add_argument(
        "-r",
        "--row",
        type=int,
        default=None,
        help="Stores data for the given row/entry of the input database \t(processes all rows otherwise)",
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="Run in verbose mode")
    parser.add_argument("--debug", action="store_true", help="Run in debug mode")
    args = parser.parse_args()

    try:
        f = open(args.inputcsv, "r")
        f.close()
    except FileNotFoundError:
        print(f"Downloading {args.inputcsv} from https://osf.io/ (CC BY 4.0 @ Geert Verdoolaege and Stanley Kaye)")
        url = "https://osf.io/zhwa3/download"
        dl = requests.get(url, allow_redirects=True)  # ,verify="/work/imas/etc/iter.pem"
        open(args.inputcsv, "wb").write(dl.content)

    if not os.path.exists(args.mapping):  # using local version
        base_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
        args.mapping = os.path.join(base_dir, "resources/mappings/h-mode-db-mapping.csv")
        if not os.path.exists(args.mapping):
            print(f"Abort! Mapping file `{args.mapping}` doesn't exists")
            exit(1)

    mf = pd.read_csv(args.mapping, keep_default_na=False, usecols=[args.varcol, args.pathcol, args.tracol])
    mf.dropna(how="all")

    if args.dbtype == "H-MODE" or args.dbtype == "HMODE":
        db = pd.read_csv(args.inputcsv, skiprows=1, keep_default_na=False, na_values=["", "-9999999", "???????"])
    elif args.dbtype == "L-MODE" or args.dbtype == "LMODE":
        db = pd.read_csv(args.inputcsv, delimiter=";", keep_default_na=False, na_values=["", "-9999999", "????????"])

    db.dropna(how="all")
    db = db.replace(to_replace=np.nan, value=None)

    rows = [args.row] if args.row is not None else db.index

    print(f"Starts mapping of the DB {args.inputcsv}, please be patient...")
    for row in rows:
        DBVAR = db.iloc[row]
        de = imas.DBEntry(
            get_backend_id(args.backend),
            args.database,
            1,
            row,
            dd_version=(args.convert_to if args.convert_to != None else args.mapping_version),
        )
        de.create()

        iod = {}
        factory = imas.IDSFactory(version=args.mapping_version)
        for ids in factory.ids_names():
            iod[ids] = factory.new(ids)
        for var in mf.loc[:, args.varcol]:
            idspath = mf[mf[args.varcol] == var].iloc[0].at[args.pathcol]

            transformation = mf[mf[args.varcol] == var].iloc[0].at[args.tracol]
            if idspath == "":
                if args.debug:
                    print(f"No mapping specified for variable {var}")
            else:
                idsname = idspath.split("/")[0]
                path = idspath.split("/")[1:]
                try:
                    ids = iod[idsname]
                except KeyError:
                    print(f"{idsname} is not a valid IDS name. Please check spelling and IDS list.")
                    sys.exit()
                if transformation != "":
                    val = evaluate(transformation)
                else:
                    val = db.at[row, var]
                if val is not None:
                    if type(val) is not str and math.isnan(val):
                        if args.verbose:
                            print(f"Nothing to store for {var} in DB entry {row}")
                    else:
                        try:
                            ids_setter(ids, path, val)
                            ids.ids_properties.homogeneous_time = 1
                        except AttributeError as ae:
                            print(ae)
                            sys.exit()

        for sids in iod.keys():
            if iod[sids].ids_properties.homogeneous_time == 1:
                if hasattr(iod[sids], "time") and iod[sids].time.size == 0:
                    iod[sids].time = iod[args.timeloc].time
                try:
                    if args.convert_to:
                        converted_ids = imas.convert_ids(iod[sids], args.convert_to)
                        de.put(converted_ids)
                    else:
                        de.put(iod[sids])
                except Exception as ex:
                    print(f"Error while attempting to write the IDS {sids} into the IMAS DB: {ex}")
                    sys.exit()
                if args.verbose:
                    print(f"The IDS {sids} was stored successfully for DB entry {row}")
        print(".", end="", flush=True)

    if args.debug:
        s = iod["summary"]
        barometry = iod["barometry"]
        print(s.boundary.type.value)
        print(s.time)
        print(s.global_quantities.ip.value)
        print(s.global_quantities.volume.value)
        print(s.elms.frequency.value)
