import os
import json
import numpy as np

from operator import itemgetter
from itertools import chain
from typing import Optional

from ..data.utils import auto_metadata_generator, get_dtypes_from_type_map, is_float  # data_harmonizer


class MetaEncoder:
    """
    A class for encoding and decoding tabular data based on its metadata, 
    including one-hot encoding for categorical variables.

    Parameters
    ----------
    col_enc : dict, optional
        A dictionary specifying custom encoding mappings for specific columns.
    **kwargs : optional
        Additional arguments for configuration.

    Attributes
    ----------
    col_enc : dict
        Custom column type definition.
    dtype_map : dict
        Mapping of column names to their data types.
    vtype_map : dict
        Mapping of column names to their variable types.
    enc_map : dict
        Mapping of column names to their encoding values.
    dec_map : dict
        Mapping of column names to their decoding values.
    num_map : dict
        Mapping of numerical column indices.
    cg_map : dict
        Mapping of categorical column indices.
    cols : list
        List of column names.
    onehot_cols : np.ndarray
        List of columns generated by one-hot encoding.
    onehot_inds : dict
        Mapping of original column names to their one-hot encoded indices.
    """
    
    def __init__(self, col_enc: Optional[dict] = None, **kwargs) -> None:
        """
        Initializes the MetaEncoder with optional column encoding constraints.

        Parameters
        ----------
        col_enc : dict, optional
            A dictionary specifying custom column types specified by the user.
        **kwargs : optional
            Additional arguments for configuration.
        """
        
        # categorical value constraints
        self.col_enc = col_enc if isinstance(col_enc, dict) else {}

        # mappers
        self.dtype_map, self.vtype_map = None, None
        self.enc_map, self.dec_map = None, None

        # continuous/categorical columns
        self.num_map = None
        self.cg_map = None

        # columns
        self.cols = None
        self.onehot_cols = None
        self.onehot_inds = None

        # index columns
        self.col_inds = None  # full column indices (list of lists)
        self.ohe_inds = None  # one-hot encoded (list of lists)
        self.or_inds = None  # non one-hot encoded (list of lists)


    def reset_params(self):
        """
        Resets all internal parameters and mappings to their initial state.

        Returns
        -------
        MetaEncoder
            The instance of the MetaEncoder with reset parameters.
        """
        
        # mappers
        self.dtype_map, self.vtype_map = None, None
        self.enc_map, self.dec_map = None, None

        # continuous/categorical columns
        self.num_map = None
        self.cg_map = None

        # columns
        self.cols = None
        self.onehot_cols = None
        self.onehot_inds = None

        # index columns
        self.ohe_inds = None  # one-hot encoded
        self.or_inds = None   # non one-hot encoded

        return self
    
    
    @staticmethod
    def is_num(x):
        """
        Checks if the input value is numerical or not.

        Parameters
        ----------
        x : object
            Input value to be checked.

        Returns
        -------
        bool
            True if the value is numerical, False otherwise.
        """
        
        try:
            float(x)
            is_num = True
        except (ValueError, TypeError):
            is_num = False
        return is_num
    
    def setup_inference(self, dtype_map: dict, onehot_inds: dict, onehot_cols: list, cols: list, enc_map: dict, dec_map: dict, **kwargs):
        # assign parameters
        self.dtype_map = dtype_map
        
        # one hot encoding
        self.onehot_inds = onehot_inds
        self.onehot_cols = np.array(onehot_cols)
        
        # mappers
        self.dec_map = {
            k1: {float(k2) if self.is_num(k2) else k2: v2 for k2, v2 in v1.items()} 
            for k1, v1 in dec_map.items()
        }
        self.enc_map = {
            k1: {float(k2) if self.is_num(k2) else k2: v2 for k2, v2 in v1.items()}
            for k1, v1 in enc_map.items()
        }

        # columns
        self.cols = cols

        # categorical columns mapper
        self.cg_map = {
            col: idx for idx, (col, map) in enumerate(self.dec_map.items()) 
            if len(map)
        }

        return self
    
    def save_params(self, dir_path: str):
        # get metadata
        meta = {
            "dtype_map": self.dtype_map,
            "onehot_inds": self.onehot_inds,
            "onehot_cols": self.onehot_cols.tolist() if self.onehot_cols is not None else None,
            "cols": self.cols,
            "dec_map": self.dec_map,
            "enc_map": self.enc_map
        }

        # save to json
        path = os.path.join(dir_path, "meta.json")
        json.dump(meta, open(path, "w"), indent=4)

        return self

    def metadata(self, data: np.ndarray, cols: Optional[list] = None, **kwargs):
        """
        Generates metadata for the provided data, including variable type mappings.

        Parameters
        ----------
        data : np.ndarray
            Input data array for which metadata is to be generated.
        cols : list, optional
            Column names. If not provided, default names are generated.
        **kwargs : optional
            Additional arguments for metadata generation.

        Returns
        -------
        dict
            A dictionary mapping column names to their inferred variable types.
        """
        
        # columns
        if not isinstance(cols, (list, np.ndarray)):
            cols = [f"attr{idx}" for idx in range(data.shape[1])]

        # get general variable mapper
        self.vtype_map = auto_metadata_generator(
            data=data,
            cols=cols,
            **kwargs,
        )

        # get new mappers
        dtype_map, opt_map = get_dtypes_from_type_map(
            vtype_map=self.vtype_map,
        )

        # assign
        self.dtype_map = dtype_map  # dtype map

        # encoder map
        self.enc_map = {
            **opt_map,
            **self.col_enc,
        }

        # decoder map
        self.dec_map = {
            k: {vv: vk for vk, vv in v.items()}
              for k, v in self.enc_map.items()
        }

        # categorical/numerical map
        num_map, cg_map = {}, {}
        for idx, (k, it) in enumerate(self.vtype_map.items()):
            # check whether it is quasi-ID
            if not it.get("type", {}).get('is_categ'):
                num_map[k] = idx
                continue

            # normal check otherwise
            if "string" in it.get("dtype"):
                cg_map[k] = idx
            elif it.get("type", {}).get("tag") == "discrete":
                cg_map[k] = idx
            else:
                num_map[k] = idx

        # assign
        self.num_map, self.cg_map = num_map, cg_map
        self.cols = list(self.dec_map.keys())

        return self.vtype_map

    @staticmethod
    def map_with_dict(
        data: np.ndarray,
        val_map: dict,
        dtype_map: Optional[dict] = {},
        column_names: Optional[list] = None,
    ):
        """
        Maps the values in the data array based on the provided value mapping.

        Parameters
        ----------
        data : np.ndarray
            Input data to be mapped.
        val_map : dict
            Mapping of column names to their encoding values.
        dtype_map : dict, optional
            Mapping of column names to their data types.
        column_names : list, optional
            Full list of column names.

        Returns
        -------
        np.ndarray
            A new array with values mapped according to the provided mapping.
        """
        
        # new data array decoded
        data_to_map = []

        # columns order
        if isinstance(column_names, list) and len(column_names) == data.shape[1]:
            cols = column_names
        else:
            cols = val_map.keys()

        # loop over mappers
        for idx, col_ in enumerate(cols):
            # map
            map_ = val_map.get(col_, {})

            # variable type mapper
            dtype = dtype_map.get(col_, "float")

            # auxiliary data
            aux_data = data[:, idx]

            if len(map_) > 0:
                # dtype encoding
                if dtype == float:
                    # map values
                    aux_val = list(
                        map(
                            lambda x: float(x) if is_float(x) else np.nan,
                            aux_data,
                        ),
                    )

                    # force float type
                    aux_val = np.array(aux_val, dtype=float)
                else:
                    # no map required
                    aux_val = data[:, idx]

                # encoded column
                aux_ser = list(
                    map(lambda x: map_.get(x, np.nan), aux_val),
                )

                # append
                data_to_map.append(aux_ser)

            else:
                data_to_map.append(data[:, idx].astype(dtype))

        # final data array
        data_to_map = np.array(data_to_map, dtype=object).swapaxes(0, 1)

        return data_to_map

    def onehot_encode(
        self,
        data: np.ndarray,
        column_names: list = None,
        **kwargs,
    ):
        """
        Applies one-hot encoding to the specified columns of the input data.

        Parameters
        ----------
        data : np.ndarray
            Input data to be one-hot encoded.
        column_names : list, optional
            Full list of column names.

        Returns
        -------
        np.ndarray
            A new array with one-hot encoded columns.
        """
        
        # full col ind map
        col_ind_map = {**self.num_map, **self.cg_map}

        # data
        if column_names is None:
            cols, cols_ind = zip(*col_ind_map)
        else:
            cols, cols_ind = column_names, itemgetter(*column_names)(col_ind_map)

        # get correct column order
        data_t = data[:, cols_ind]

        # new column map
        onehot_cols, data_onehot = [], []
        onehot_inds = {}

        counter = 0  # count indices
        for idx, col in enumerate(cols):
            # auxiliary indices
            aux_ind = []

            # get options
            aux_opt = self.enc_map.get(col)

            # if not empty or non-binary (categories are present)
            if len(aux_opt) > 2 and col in self.cg_map:
                aux_col, aux_data = [], []
                for opt in aux_opt.values():
                    # new col
                    aux_col += [f"{col}_OH-{opt}"]

                    # new onehot encoded data
                    aux_data += [np.where(data_t[:, idx] == opt, 1, 0)]

                    # append
                    aux_ind.append(counter)

                    # increment
                    counter += 1
            else:
                # original col
                aux_col = [f"{col}_OR"]

                # original data
                aux_data = [data[:, idx]]

                # append
                aux_ind.append(counter)

                # increment
                counter += 1

            # append
            onehot_cols += aux_col  # columns
            onehot_inds[col] = aux_ind  # col: indices
            data_onehot += aux_data  # data

        # assign
        if self.onehot_cols is None:
            self.onehot_cols = np.array(onehot_cols)
            self.onehot_inds = onehot_inds

            # final column indices
            self.col_inds = list(onehot_inds.values())
            self.ohe_inds = list(chain(*[inds for _, inds in onehot_inds.items() if len(inds) > 1]))
            self.or_inds = list(chain(*[inds for _, inds in onehot_inds.items() if len(inds) == 1]))

        # reshape data
        data_onehot = np.array(data_onehot).swapaxes(0, 1)

        return data_onehot

    def onehot_decode(
        self,
        data: np.ndarray,
        **kwargs,
    ):
        """
        Decodes one-hot encoded data back to its original representation.

        Parameters
        ----------
        data : np.ndarray
            One-hot encoded data to be decoded.

        Returns
        -------
        np.ndarray
            A new array with decoded values.
        """
        
        # new column map
        data_inv = []

        for col, inds in self.onehot_inds.items():
            # column has been onehot encoded
            if len(inds) > 1:
                # get encoding values
                values = np.array([int(oh_col.split("_")[-1].split("-")[-1]) for oh_col in self.onehot_cols[inds]])

                # curate combination vectors
                aux_cur = (data[:, inds] == data[:, inds].max(axis=1)[:, None]).astype(int)

                # new onehot decoded data
                aux_data = values[np.argmax(aux_cur, axis=1)].astype(int)

                # append
                data_inv.append(aux_data)

            # column has not been onehor encoded
            else:
                # decoding map
                aux_dec_map = self.dec_map.get(col, {})
                
                # checks (binary and categorical)
                bin_col_check = len(aux_dec_map) == 2
                cg_col_check = col in self.cg_map

                # original data (no decoding only if binary target)
                if cg_col_check and not bin_col_check:  # quasi-unique column
                    aux_data = data[:, inds[0]].round(0).astype(int)
                elif bin_col_check:  # binary categorical column
                    # bins
                    bins = list(aux_dec_map.keys())

                    # threshold
                    thresh = bins[0] + abs(bins[1] - bins[0]) / 2
                    
                    # get data
                    aux_data = np.where(data[:, inds[0]] > thresh, bins[1], bins[0])
                else:  # continuous column
                    aux_data = data[:, inds[0]].astype(float)

                # append
                data_inv.append(aux_data)

        # reshape data
        data_decoded = np.array(data_inv).swapaxes(0, 1)

        return data_decoded

    def encode(
        self,
        data: np.ndarray,
        column_names: Optional[list] = None,
        with_onehot: Optional[bool] = False,
        **kwargs,
    ):
        """
        Encodes the input data using the defined mappings and optional one-hot encoding.

        Parameters
        ----------
        data : np.ndarray
            Input data to be encoded.
        column_names : list, optional
            Full list of column names.
        with_onehot : bool, optional
            Whether to apply one-hot encoding.

        Returns
        -------
        np.ndarray
            The encoded data array.
        """
        
        # encoded (string to num)
        new_data = self.map_with_dict(
            data=data,
            val_map=self.enc_map,
            dtype_map=self.dtype_map,
            column_names=column_names,
        )

        # one hot encoding
        if with_onehot:
            new_data = self.onehot_encode(new_data, column_names)

        return new_data

    def decode(
        self,
        data: np.ndarray,
        column_names: Optional[list] = None,
        with_onehot: Optional[bool] = True,
        **kwargs,
    ):
        """
        Decodes the input data back to its original representation using defined mappings.

        Parameters
        ----------
        data : np.ndarray
            Encoded data to be decoded.
        column_names : list, optional
            Full list of column names.
        with_onehot : bool, optional
            Whether to decode one-hot encoded columns. True if one-hot encoding was performed, False otherwise.

        Returns
        -------
        np.ndarray
            The decoded data array.
        """
        
        # one hot decoding
        new_data = self.onehot_decode(data) if with_onehot else data.copy()

        # remap to original values
        new_data = self.map_with_dict(
            data=new_data,
            val_map=self.dec_map,
            dtype_map=self.dtype_map,
            column_names=column_names,
        )

        return new_data
