# -*- coding: utf-8 -*-
"""
Created on Tue Aug 24 22:15:00 2024

@author: yqlim

Wrappers for xgboost.
"""
import logging
logger = logging.getLogger(__name__)

from typing import Callable
import os
import re
import pandas as pd
import xgboost  

class IteratorCsv(xgboost.DataIter):
    # https://xgboost.readthedocs.io/en/stable/tutorials/external_memory.html
    # https://stackoverflow.com/questions/78348470/xgboost-aft-survival-model-with-external-memory-iterator
    def __init__(self, 
            filepaths: list[str],
            y_label: str,
            x_labels: list[str] = None,
            x_labels_exact_match: bool = False            
        ) -> None:
        self._file_paths = filepaths
        self.y_label = y_label
        self.x_labels = x_labels
        self.x_labels_exact_match = x_labels_exact_match

        self.x_labels_final = None if not self.x_labels_exact_match else x_labels

        self._it = 0
        # XGBoost will generate some cache files under current directory
        # with the prefix "cache"
        super().__init__(cache_prefix=os.path.join('.', 'cache'))

    def _preprocess(self, df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
        """
        Split DataFrame into X, y.

        Parameters
        ----------
        df : pd.DataFrame
            from loading csv
        """
        if not self.x_labels_exact_match:
            regex = re.compile(
                '|'.join(self.x_labels)
            )
            self.x_labels_final = [
                col for col in df.columns
                if regex.search(col)
            ]

        df = df.fillna(-1)
        if not self.x_labels_final is None:
            X = df[self.x_labels_final]
        else:
            X = df.drop(self.y_label, axis=1)

        return (X, df[self.y_label].astype('float32'))

    def next(self, input_data: Callable) -> int:
        """
        Advance the iterator by 1 step and pass the data to XGBoost.
        This function is called by XGBoost during the construction
        of ``DMatrix``
        """
        if self._it == len(self._file_paths):
            # return 0 to let XGBoost know this is the end of iteration
            return 0
        
        # input_data is a function passed in by XGBoost who has the 
        # exact same signature of ``DMatrix``
        df = pd.read_csv(self._file_paths[self._it])
        while df.empty:
            self._it += 1
            if self._it == len(self._file_paths):
                return 0
            df = pd.read_csv(self._file_paths[self._it])

        X, y = self._preprocess(df)
        
        print(f"X: {X.shape}; y: {y.shape}")
        print(f"X: {X.isna().values.any()}; y: {y.isna().values.any()}")
        input_data(data=X, label=y)
        self._it += 1
        # Return 1 to let XGBoost know we have not seen all the files yet
        return 1

    def reset(self) -> None:
        """
        Reset iterator to beginning.
        """
        self._it = 0

