"""Data processing utils for MIDAS2."""

import pandas as pd
import numpy as np


def _process_missing(X):
    """
    Set missings to 0 and return missing matrix.

    """
    data = X.copy()
    data = np.nan_to_num(data, 0)
    missings = np.isnan(X)
    return data, missings


def _format_cols(data: pd.DataFrame):
    """
    Format input pandas DataFrame for use in MIDAS model.

    Parameters:
    data: pd.DataFrame
    verbose: bool. Print column information post-transformation.
        This is useful for checking the auto-conversion has worked.


    """

    data = data.copy()

    # hash maps for column types
    type_dict = {}

    # sanity check
    if not isinstance(data, pd.DataFrame):
        raise ValueError("Data must be a pandas DataFrame")

    # store orderings
    col_names = []
    col_types = []

    for col in data.columns:
        if data[col].dtype == "object":
            data[col] = data[col].astype("category")

        if data[col].dtype == "boolean":
            data[col] = 1 * data[col]
            col_types.append("bin")
            col_names += [col]
            type_dict[col] = [False, True]

        elif data[col].dtype.name == "category":

            cats = data[col].cat.categories
            na_tmp = data[col].isnull()

            # handle binary categorical columns
            if len(cats) == 2:
                tmp_col = 1 * data[col].eq(cats[1])
                tmp_col[na_tmp] = np.nan
                data[col] = tmp_col
                col_types.append("bin")
                col_names += [col]
            else:
                tmp_df = pd.get_dummies(data[col], prefix=col, dtype=float)
                tmp_df[na_tmp] = np.nan
                data = data.drop(col, axis=1, inplace=False)
                data = pd.concat([data, tmp_df], axis=1)
                col_types.append(tmp_df.shape[1])
                col_names += list(tmp_df.columns.values)

            type_dict[col] = cats

        elif data[col].dtype.name[:5] == "float" or data[col].dtype.name[:3] == "int":

            unq_vals = data[col].astype(float).unique()
            unq_vals = unq_vals[~np.isnan(unq_vals)]  # remove NaNs for comparison

            # handle binary columns
            if set(unq_vals) == {0.0, 1.0}:
                type_dict[col] = [0.0, 1.0]
                col_types.append("bin")
            elif data[col].min() >= 0:
                type_dict[col] = "positive"
                col_types.append("pos")
            else:
                type_dict[col] = "numeric"
                col_types.append("num")

            col_names += [col]
        else:
            raise ValueError(
                f"Column {col} has unrecognized type. Please convert to numeric or categorical."
            )

    return [data.loc[:, col_names], type_dict, col_types]


def _format_cols_test(data: pd.DataFrame, col_types, type_dict, col_names):
    """
    Format input pandas DataFrame for use in MIDAS model.

    Parameters:
    data: pd.DataFrame
    verbose: bool. Print column information post-transformation.
        This is useful for checking the auto-conversion has worked.


    """
    # sanity check
    if not isinstance(data, pd.DataFrame):
        raise ValueError("Data must be a pandas DataFrame")

    data = data.copy()

    for col in data.columns:
        if data[col].dtype == "object":
            data[col] = data[col].astype("category")

        if data[col].dtype == "boolean":
            data[col] = 1 * data[col]

        elif data[col].dtype.name == "category":

            cats = type_dict[col]
            na_tmp = data[col].isnull()
            idx = data[col].map(
                lambda v: cats.get_loc(v) if (pd.notna(v) and v in cats) else np.nan
            )

            if len(cats) == 2:
                tmp_col = (idx == 1).astype(float)
                tmp_col[idx.isna()] = np.nan
                data[col] = tmp_col
            else:

                cats_df = pd.DataFrame(
                    0.0, index=data.index, columns=[f"{col}_{c}" for c in cats]
                )

                for j, cat_name in enumerate(cats):
                    cats_df.loc[idx == j, f"{col}_{cat_name}"] = 1.0
                cats_df.loc[idx.isna(), :] = np.nan
                data = data.drop(col, axis=1).join(cats_df)

        elif data[col].dtype.name[:5] == "float" or data[col].dtype.name[:3] == "int":
            pass

        else:
            raise ValueError(
                f"Column {col} has unrecognized type. Please convert to numeric or categorical."
            )

    return data.loc[:, col_names]
