# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['BusinessRule']

from typing import Union, List, Dict, Tuple
from pathlib import Path

import numpy as np
import pandas as pd

from sklearn.base import BaseEstimator
from sklearn.metrics import accuracy_score, mean_squared_error

from igraph import Graph

from .storable import Storable


def generate_range_mask(range_dict:dict, X:pd.DataFrame, kind:str='all')->pd.Series:
    """generates boolean mask for X based on range_dict dictionary. 

        range_dict should be of the format 
        ```
        range_dict = {
            'petal length (cm)': [[4.1, 4.7], [5.2, 7.5]],
            'petal width (cm)': [1.6, 2.6]
        }
        ```
        
        Can also be categorical:
        
        ```
        range_dict = {
            'gender': ['male']
        }
        ```
        
    Args:
        range_dict (dict): dictionary describing which ranges for which columns should be evaluated
        X (pd.DataFrame): input dataframe
        kind (str, optional): Only return True if range conditions hold for all cols ('all') 
            or if range conditions hold for any column ('any'). Defaults to 'all'.

    Returns:
        pd.Series: boolean mask
    """
    
    def get_mask(X, col, col_range):
        if isinstance(col_range[0], str):
            return X[col].isin(col_range)
        else:
            return (X[col] > col_range[0]) & (X[col] < col_range[1])

    def AND_masks(masks):
        for i, mask in enumerate(masks):
            combined_mask = mask if i == 0 else combined_mask & mask # noqa: F82
        return combined_mask

    def OR_masks(masks):
        for i, mask in enumerate(masks):
            combined_mask = mask if i == 0 else combined_mask | mask # noqa: F82
        return combined_mask

    def generate_range_mask(X, col, col_range):
        if isinstance(col_range[0], list):
            return OR_masks([get_mask(X, col, cr) for cr in col_range])
        return get_mask(X, col, col_range)
    
    if not range_dict:
        return pd.Series(np.full(len(X), False), index=X.index)
    if kind == 'all':
        return AND_masks([generate_range_mask(X, col, col_range) for col, col_range in range_dict.items()])
    elif kind == 'any':
        return OR_masks([generate_range_mask(X, col, col_range) for col, col_range in range_dict.items()])
    else:
        raise ValueError("ValueError! Only kind='all' and kind='any' are supported!")


class BusinessRule(BaseEstimator, Storable):
    def __init__(self, prediction=None, default=None):
        self._store_child_params(level=2)
        if not hasattr(self, "prediction"):
            self.prediction = prediction
        if not hasattr(self, "default"):
            self.default = default
            
        if self.prediction is None:
            self.prediction = np.nan
        if self.default is None:
            self.default = np.nan
            
        self._rule_id = None
        
    def fit(self, X:pd.DataFrame, y:Union[pd.Series, np.ndarray]=None):
        pass

    def predict(self, X:pd.DataFrame)->np.ndarray:
        assert hasattr(self, "__rule__"), "You need to implement the __rule__ method first!"
        return np.where(self.__rule__(X), self.prediction, self.default)
    
    def _score_rule(self, y, y_preds, mask, prediction, default, 
                    scores_df=None, is_classifier=False)->pd.DataFrame:
        if scores_df is None:
            scores_df = pd.DataFrame(columns=[
                'rule_id', 'name','description', 'prediction', 
                'n_inputs', 'n_outputs','coverage', 
                'accuracy' if is_classifier else 'rmse'
            ])
        score_dict = dict(
                rule_id=self._rule_id, name=self.__class__.__name__,
                description=self.__rulerepr__(),
                prediction = prediction,
                n_inputs=len(y), n_outputs=mask.sum(), 
                coverage = mask.mean() if len(mask)>0 else np.nan,
        )
        
        if is_classifier:
            if len(y[mask]) > 0:
                score_dict['accuracy'] = accuracy_score(y[mask], y_preds[mask])
            else:
                score_dict['accuracy'] = np.nan
        else:
            if len(y[mask]) > 0:
                score_dict['rmse'] = mean_squared_error(y[mask], y_preds[mask], squared=False)
            else:
                score_dict['rmse'] = np.nan
        
        scores_df = scores_df.append(score_dict, ignore_index=True)
        
        if not np.isnan(default):
            default_score_dict = dict(
                    rule_id=self._rule_id, 
                    name=chr(int("21B3", 16)), #self.__class__.__name__,
                    description=f"default: predict {self.default}",
                    prediction=default,
                    n_inputs=len(y), n_outputs=np.invert(mask).sum(), 
                    coverage = np.invert(mask).mean() if len(mask)>0 else np.nan,       
            )
            if is_classifier:
                if np.invert(mask).sum() > 0:
                    default_score_dict['accuracy'] = accuracy_score(
                        y[~mask], np.full(np.invert(mask).sum(), default))
                else:
                    default_score_dict['accuracy'] = np.nan      
            else:
                if np.invert(mask).sum() > 0:
                    default_score_dict['rmse'] = mean_squared_error(
                        y[~mask], np.full(np.invert(mask).sum(), default), squared=False)
                else:
                    default_score_dict['rmse'] = np.nan
                    
            scores_df = scores_df.append(default_score_dict, ignore_index=True)
        return scores_df
    
    def score_rule(self, X:pd.DataFrame, y:Union[np.ndarray, pd.Series], 
                   scores_df:pd.DataFrame=None, is_classifier:bool=False)->pd.DataFrame:
        mask = pd.Series(self.__rule__(X)).values
        y_preds = self.predict(X)
        return self._score_rule(y, y_preds, mask, 
                                self.prediction, self.default, 
                                scores_df, is_classifier)
    
    def set_rule_id(self, rule_id:int=0)->int:
        self._rule_id = rule_id
        return rule_id+1
    
    def get_max_rule_id(self, max_rule_id:int=0)->int:
        if self._rule_id is not None and self._rule_id > max_rule_id:
            return self._rule_id
        return max_rule_id
    
    def get_rule(self, rule_id:int):
        if self._rule_id is not None and self._rule_id == rule_id:
            return self
        
    def replace_rule(self, rule_id:int, new_rule)->None:
        if self._rule_id is not None and self._rule_id == rule_id:
            assert isinstance(new_rule, BusinessRule)
            self.__class__ = new_rule.__class__
            self.__dict__ = new_rule.__dict__
            return self

    def remove_rule(self, rule_id:int):
        return None
            
    def get_rule_params(self, rule_id:int)->dict:
        if self._rule_id is not None and self._rule_id == rule_id:
            return self.get_params()
    
    def set_rule_params(self, rule_id:int, **params)->None:
        if self._rule_id is not None and self._rule_id == rule_id:
            self.set_params(**params)
          
    def get_rule_input(self, rule_id:int, X:pd.DataFrame, y:Union[pd.Series, np.ndarray]=None
                      )->Union[pd.DataFrame, Tuple[pd.DataFrame, Union[pd.Series, np.ndarray]]]:
        if self._rule_id is not None and self._rule_id == rule_id:
            if y is not None:
                return X, y
            else:
                return X
            
        if y is not None:
            return None, None
        else:
            return None
        
    def get_rule_leftover(self, rule_id:int, X:pd.DataFrame, y:Union[pd.Series, np.ndarray]=None
                         )->Union[pd.DataFrame, Tuple[pd.DataFrame, Union[pd.Series, np.ndarray]]]:
        if self._rule_id is not None and self._rule_id == rule_id:
            mask = np.invert(pd.Series(self.__rule__(X)).values)
            if y is not None:
                return X[mask], y[mask]
            else:
                return X[mask]
            
        if y is not None:
            return None, None
        else:
            return None
                    
    def get_params(self, deep:bool=True)->dict:
        """ """
        return self._stored_params

    def set_params(self, **params)->None:
        """ """
        for k, v in params.items():
            if k in self._stored_params:
                self._stored_params[k] = v
                setattr(self, k, v)

    def _get_casewhens(self, casewhens:dict=None):
        if casewhens is None:
            return {}
        else:
            return casewhens

    def _get_binarynodes(self, binarynodes:dict=None):
        if binarynodes is None:
            return {}
        else:
            return binarynodes
                
    def add_to_igraph(self, graph:Graph=None)->Graph:
        if graph is None:
            graph = Graph()
            graph.vs.set_attribute_values('rule_id', [])
            graph.vs.set_attribute_values('name', [])
            graph.vs.set_attribute_values('description', [])
            graph.vs.set_attribute_values('rule', [])
            graph.es.set_attribute_values('casewhen', [])
            graph.es.set_attribute_values('binary_node', [])
        graph.add_vertex(
            rule_id=self._rule_id,
            name=self.__class__.__name__,
            description=self.__rulerepr__(),
            rule=self
        )
        return graph
        
    def __rulerepr__(self)->str:
        return "BusinessRule"

    def to_yaml(self, filepath:Union[Path, str]=None, return_dict:bool=False):
        """Store object to a yaml format.

        Args:
            filepath: file where to store the .yaml file. If None then just return the
                yaml as a str.
            return_dict: instead of return a yaml str, return the raw dict.

        """
        return super().to_yaml(filepath, return_dict, comment=self.describe())



