# -*- coding: utf-8 -*-
"""Created on Fri Jul  6 13:52:04 2018.

@author: erwan
"""

# TODO refactor : rename this file as hitran_utils.py

from warnings import warn

import numpy as np
import pandas as pd

from ..misc.utils import NotInstalled, not_installed_vaex_args

try:
    import vaex
except ImportError:
    vaex = NotInstalled(*not_installed_vaex_args)


from ..misc.warning import PerformanceWarning


def parse_hitran_file(fname, columns, count=-1, output="pandas", molecule=None):
    """Parse a file under HITRAN ``par`` format. Parsing is done in binary
    format with :py:func:`numpy.fromfile` so it's as fast as possible.

    Parameters
    ----------
    fname: str
        filename.
    columns: dict
        list of columns and their format.

    Other Parameters
    ----------------
    count: int
        number of lines to read. If ``-1`` reads all file.

    output : str
        specifies the output type

    Returns
    -------
    df: pandas DataFrame or Vaex Dataframe
        dataframe with lines.

    See Also
    --------
    Used in :py:func:`~radis.api.hitranapi.hit2df` and :py:func:`~radis.api.cdsdapi.cdsd2df`
    """

    # To be faster, we read file totally in bytes mode with fromfiles. But that
    # requires to properly decode the line return character:

    # We cannot simply infer it from the OS :
    # problem arise when file was written in an OS and read in another OS (for instance,
    # line return characters are not converted when read from .egg files). Here
    # we read the first line and infer the line return character for it
    data = _read_hitran_file(
        fname, columns, count=1, linereturnformat="a2"
    )  # 'a2' allocates space to get \n or \n\r linereturn formats
    linereturnformat = _get_linereturnformat(data, columns, fname)

    # Now re-read with correct line return character
    data = _read_hitran_file(fname, columns, count, linereturnformat)

    # Return a Pandas dataframe
    df = _ndarray2df(data, columns, linereturnformat, molecule=molecule)

    if output == "vaex":
        if isinstance(vaex, NotInstalled):
            raise ImportError(
                "Vaex is not available. Please install vaex or use a different engine."
            )
        df = vaex.from_pandas(df)
    return df


def _get_linereturnformat(data, columns, fname=""):
    """
    Get line return character & format (size).

    Notes
    -----

    We cannot simply infer it from the OS :
    problem arise when file was written in an OS and read in another OS (for instance,
    line return characters are not converted when read from .egg files). Here
    we read the first line and infer the line return character for it
    """
    # fname just for the error message

    # get format (size) of line return
    from radis.misc.basics import to_str

    linereturn = to_str(data[0][-1])
    if to_str("\r\n") in linereturn:
        linereturnformat = "a2"
    elif to_str("\n") in linereturn or to_str("\r") in linereturn:
        linereturnformat = "a1"
    else:
        raise ValueError(
            f"Unknown Line return format: {linereturn}. Check that your file {fname} has the HITRAN format. First line : {data[0]}"
        )

    return linereturnformat


def _ndarray2df(data, columns, linereturnformat, molecule=None):
    """ """

    # ... Cast to new type
    # This requires to recast all the data already read, but is still the fastest
    # method I found to read a file directly (for performance benchmark see
    # CDSD-HITEMP parser)
    newtype = [c[0] if (c[1] == str) else c[1] for c in columns.values()]
    dtype = list(zip(list(columns.keys()), newtype)) + [
        ("_linereturn", linereturnformat)
    ]

    # Override iso type if CO2
    if molecule == "CO2":
        dtype[1] = (
            "iso",
            "<U1",
        )  # in HITEMP2024, isotopologue 10, 11, 12 are 0, A, B. Converted later
    data = _cast_to_dtype(data, dtype)

    # %% Create dataframe
    df = pd.DataFrame(data, columns=list(columns.keys()) + ["_linereturn"])

    # Delete dummy column than handled the line return character
    del df["_linereturn"]

    # Update format
    for k, c in columns.items():
        if c[1] == str:
            df[k] = df[k].str.decode("utf-8")

    # Replace CO2 isotopologue labels with numeric codes
    if molecule == "CO2":
        df["iso"] = df["iso"].replace({"0": 10, "A": 11, "B": 12}).astype(int)

    # Strip whitespaces around PQR columns (due to 2 columns jumped)
    if "branch" in df:  # (only in CDSD)
        df["branch"] = df.branch.str.strip()

    return df


def _read_hitran_file(fname, columns, count, linereturnformat):
    """
    Returns
    -------

    data: np.ndarray
    """

    dt = _create_dtype(columns, linereturnformat)

    return np.fromfile(fname, dtype=dt, count=count)


def _create_dtype(columns, linereturnformat):
    """Create dtype from specific columns.

    Returns
    -------
    dt: dtypes
    """
    # ... Create a dtype with the binary data format and the desired column names
    dtype = [(k, c[0]) for (k, c) in columns.items()] + [
        ("_linereturn", linereturnformat)
    ]
    # ... _linereturn is to capture the line return symbol. We delete it afterwards
    return _format_dtype(dtype)


def _format_dtype(dtype):
    """Format dtype from specific columns.

    Crash with hopefully helping error message
    """
    try:
        dt = np.dtype(
            [(str(k), c) for k, c in dtype]
        )  # DeprecationWarning: Data type alias 'a' was deprecated in NumPy 2.0. Use the 'S' alias instead.
        # Note: dtype names cannot be `unicode` in Python2. Hence the str()
    except TypeError as err:
        # Cant read database. Try to be more explicit for user before crashing
        try:
            print("Data type")
            print("-" * 30)
            for k, c in dtype:
                print(str(k), "\t\t", c)
            print("-" * 30)
        finally:
            raise TypeError("Could not read datatype. See details above.") from err
    return dt


def _cast_to_dtype(data, dtype):
    """Cast array to certain type, crash with hopefully helping error message.

    Return casted data.

    Parameters
    ----------
    data: array to cast
    dtype: (ordered) list of (param, type)
    """

    dt = _format_dtype(dtype)
    try:
        data = np.array(data, dtype=dt)
    except ValueError as err:
        try:
            # Cant read database. Try to be more explicit for user before crashing
            # ... identify faulty row
            print("Cant cast data to specific dtype. Trying row by row:")
            for r in range(len(data)):
                try:
                    np.array(data[r], dtype=dt)
                except ValueError:
                    break
            print(f"Error may be on row {r}:")
            print("-" * 30)
            for i in range(len(data[r])):
                print(i, dtype[i], "\t\t", data[r][i])
                print(i, dtype[i], "\t\t", np.array(data[r][i], dtype=dt[i]))
            print("-" * 30)
            print(">>> Next param:", dtype[i], ". Value:", data[r][i], "\n")
        finally:
            raise ValueError(
                "Cant cast data to specific dtype. Tried column by column. See results above"
            ) from err

    return data


def drop_object_format_columns(df, verbose=True):
    """Remove 'object' columns in a pandas DataFrame. If Vaex Dataframe, there is by construction no object in the Dataframe

    They are not useful to us at this time, and they slow down all
    operations (as they are converted to 'object' in pandas DataFrame).
    If you want to keep them, better convert them to some numeric values
    """
    # Remove ["ierr", "iref", "lmix", "Fu", "syml", "Fl"] which is object for Python<=3.12 and string for Python 3.13+
    additional_cols = ["ierr", "iref", "lmix", "Fu", "syml", "Fl"]
    df_type = type(df)
    objects = [k for k, v in df.dtypes.items() if v == object]

    cols_to_drop = list(
        set([col for col in additional_cols if col in df.columns] + objects)
    )
    if df_type == pd.DataFrame:
        df.drop(cols_to_drop, axis=1, inplace=True)
    elif (
        not isinstance(vaex, NotInstalled) and df_type == vaex.dataframe.DataFrameLocal
    ):  # no objects in vaex
        df.drop(cols_to_drop, inplace=True)
    else:
        raise NotImplementedError(df_type)

    if verbose >= 2 and len(objects) > 0:
        print(
            (
                f"The following columns had the `object` format and were removed: {objects}. "
                f"The following additional columns were also removed: {cols_to_drop}"
            )
        )
    return df


def replace_PQR_with_m101(df):
    """Return O, P, Q, R, S in column ``branch`` with -2, -1, 0, 1, 2 to get a fully numeric
    database. This improves performances quite a lot, as Pandas doesnt have a
    fixed-string dtype hence would use the slow ``object`` dtype.

    Parameters
    ----------

    df: pandas Dataframe or Vaex Dataframe
        ``branch`` must be a column name.


    Returns
    -------

    None:
        ``df`` is is modified in place
    """

    # Note: somehow pandas updates dtype automatically. We have to check
    # We just have to replace the column:

    # First checking if the input type is correct
    df_type = type(df)
    if df_type == pd.DataFrame:
        pass
    elif (
        not isinstance(vaex, NotInstalled) and df_type == vaex.dataframe.DataFrameLocal
    ):
        pass
    else:
        raise NotImplementedError(df_type)

    # Then do the actual mapping
    if df.dtypes["branch"] != np.float64:
        mapping = {
            "P": -1,
            "Q": 0,
            "R": 1,
            "O": -2,
            "S": 2,
            # Also handle string numeric values for fast_parsing
            "0": 0,
            "1": 1,
            "2": 2,
            "-1": -1,
            "-2": -2,
        }
        valid_values = [-2, -1, 0, 1, 2]

        if df_type == pd.DataFrame:  # pandas
            pd.set_option("future.no_silent_downcasting", True)
            new_col = df["branch"].replace(mapping).infer_objects(copy=False)
            pd.set_option("future.no_silent_downcasting", False)

            invalid_values = new_col[~new_col.isin(valid_values)]
            # Raise a warning if there are branch values outside valid_values
            if not invalid_values.empty:
                invalid_indices = invalid_values.index
                warn(
                    f"Invalid values found in 'branch' column.\n"
                    f"The following branch values are not within [-2, -1, 0, 1, 2] (only 5 first index out {len(invalid_indices)} of are shown): \n"
                    f"{df.loc[invalid_indices[:5], :]}",
                    UserWarning,
                )
        else:  # vaex
            new_col = df["branch"].map(mapping, allow_missing=True)

            # invalid_values = df["branch"][~new_col.isin(valid_values)] #not implemented in vaex
            is_invalid = ~new_col.isin(valid_values)
            # Raise a warning if there are branch values outside valid_values
            if is_invalid.unique() != [False]:
                invalid_indices = invalid_values.index
                warn(
                    "Invalid values found in 'branch' column.\nValid values are branch = [-2, -1, 0, 1, 2]",
                    UserWarning,
                )

        new_col = new_col.astype("float")
        df["branch"] = new_col

    if df.dtypes["branch"] not in (int, float):
        warn(
            message=(
                f"Expected branch data type to be 'int64', "
                f"got '{df.dtypes['branch']}' instead."
                f"This warning is safe to ignore although it is going to "
                f"reduce the performance of your calculations. "
                f"For further details, see:"
                f"https://github.com/radis/radis/issues/65"
            ),
            category=PerformanceWarning,
        )