# -*- coding: utf-8 -*-
import brifc
import pandas as pd
import numpy as np
from pandas.api import types as ptypes


def _infer_var_type(dtype):
    if ptypes.is_float_dtype(dtype):
        return "numeric"
    if ptypes.is_integer_dtype(dtype) and not ptypes.is_bool_dtype(dtype):
        return "integer"
    if ptypes.is_bool_dtype(dtype) or ptypes.is_categorical_dtype(dtype) or ptypes.is_object_dtype(dtype) or ptypes.is_string_dtype(dtype):
        return "factor"
    raise ValueError("Column types must be numeric, integer or factor-like (object/string/category/bool).")


def _collect_indices(coltypes, offset=0):
    vartypes = ["" for _ in range(len(coltypes))]
    num_ind = []
    int_ind = []
    fac_ind = []
    for j, dtype in enumerate(coltypes):
        var_type = _infer_var_type(dtype)
        vartypes[j] = var_type
        if var_type == "numeric":
            num_ind.append(j + offset)
        elif var_type == "integer":
            int_ind.append(j + offset)
        else:
            fac_ind.append(j + offset)
    return vartypes, num_ind, int_ind, fac_ind

class brif:
    tmp_preddata = "tmp_brif_preddata.txt"
    max_factor_levels = 30
    max_integer_classes = 20
    n_numeric_cuts = 31
    n_integer_cuts = 31
    ps = 0
    max_depth = 20
    min_node_size = 1
    ntrees = 200
    nthreads = 4
    bagging_method = 0
    bagging_proportion = 0.9
    split_search = 4
    search_radius = 5
    vote_method = 1
    seed = 2022
    na_numeric = 0
    na_integer = -9999
    na_factor = 'NA'
    _model = None
    _problem_type = 'classification'
    def __init__(self, param = dict()):
        if "tmp_preddata" in param:
            self.tmp_preddata = param['tmp_preddata']
        if "max_factor_levels" in param:
            self.max_factor_levels = param['max_factor_levels']
        if "max_integer_classes" in param:
            self.max_integer_classes = param['max_integer_classes']
        if "n_numeric_cuts" in param:
            self.n_numeric_cuts = param['n_numeric_cuts']
        if "n_integer_cuts" in param:
            self.n_integer_cuts = param['n_integer_cuts']        
        if "ps" in param:
            self.ps = param['ps']
        if "max_depth" in param:
            self.max_depth = param['max_depth']   
        if "min_node_size" in param:
            self.min_node_size = param['min_node_size']
        if "ntrees" in param:
            self.ntrees = param['ntrees']   
        if "nthreads" in param:
            self.nthreads = param['nthreads']
        if "bagging_method" in param:
            self.bagging_method = param['bagging_method']   
        if "bagging_proportion" in param:
            self.bagging_proportion = param['bagging_proportion']
        if "split_search" in param:
            self.split_search = param['split_search']   
        if "search_radius" in param:
            self.search_radius = param['search_radius']
        if "vote_method" in param:
            self.vote_method = param['vote_method']    
        if "seed" in param:
            self.seed = param['seed']
    def set_param(self, param = dict()):
        if "tmp_preddata" in param:
            self.tmp_preddata = param['tmp_preddata']
        if "max_factor_levels" in param:
            self.max_factor_levels = param['max_factor_levels']
        if "max_integer_classes" in param:
            self.max_integer_classes = param['max_integer_classes']
        if "n_numeric_cuts" in param:
            self.n_numeric_cuts = param['n_numeric_cuts']
        if "n_integer_cuts" in param:
            self.n_integer_cuts = param['n_integer_cuts']        
        if "ps" in param:
            self.ps = param['ps']
        if "max_depth" in param:
            self.max_depth = param['max_depth']   
        if "min_node_size" in param:
            self.min_node_size = param['min_node_size']
        if "ntrees" in param:
            self.ntrees = param['ntrees']   
        if "nthreads" in param:
            self.nthreads = param['nthreads']
        if "bagging_method" in param:
            self.bagging_method = param['bagging_method']   
        if "bagging_proportion" in param:
            self.bagging_proportion = param['bagging_proportion']
        if "split_search" in param:
            self.split_search = param['split_search']   
        if "search_radius" in param:
            self.search_radius = param['search_radius']
        if "vote_method" in param:
            self.vote_method = param['vote_method']         
        if "seed" in param:
            self.seed = param['seed']  
        if "na_numeric" in param:
            self.na_numeric = param['na_numeric']   
        if "na_integer" in param:
            self.na_integer = param['na_integer']
        if "na_factor" in param:
            self.na_factor = param['na_factor'] 
            
    def get_param(self):
        return(dict({
                    "tmp_preddata":self.tmp_preddata,
                    "max_factor_levels":self.max_factor_levels,
                    "max_integer_classes":self.max_integer_classes,
                    "n_numeric_cuts":self.n_numeric_cuts,
                    "n_integer_cuts":self.n_integer_cuts,
                    "ps":self.ps,
                    "max_depth":self.max_depth,
                    "min_node_size":self.min_node_size,
                    "ntrees":self.ntrees,
                    "nthreads":self.nthreads,
                    "bagging_method":self.bagging_method,
                    "bagging_proportion":self.bagging_proportion,
                    "split_search":self.split_search,
                    "search_radius":self.search_radius,
                    "vote_method":self.vote_method,
                    "seed":self.seed,
                    "na_numeric":self.na_numeric,
                    "na_integer":self.na_integer,
                    "na_factor":self.na_factor
                    }))
    
    def fit(self, df, target_col):
        if not isinstance(df, pd.DataFrame):
            raise ValueError("df must be a pandas data frame.")
        n = df.shape[0]
        if n < 16:
            raise ValueError("Too few training data points. Model cannot run.")
        p = df.shape[1] - 1
        if p < 1:
            raise ValueError("Too few predictors.")
        target_col_num = 0
        colnames = list(df.columns.values)
        if isinstance(target_col, str):
            target_col_num = colnames.index(target_col)
        elif isinstance(target_col, int):
            target_col_num = target_col
        else:
            raise ValueError('target_col must be str or int.')
        if target_col_num < 0 or target_col_num > p:
            raise ValueError('target_col_num out of bound.')
        
        target_col_name = colnames[target_col_num]
        del colnames[target_col_num]
        colnames = [target_col_name] + colnames
        df = df[colnames]
        coltypes = df.dtypes
        vartypes, num_ind, int_ind, fac_ind = _collect_indices(coltypes, offset=0)
        
        # check unique levels
        for j in fac_ind:
            nu = len(df.iloc[:,j].unique())
            if(nu > self.max_factor_levels):
                raise ValueError("Variable " + str(colnames[j]) + ' has ' + str(nu) + " unique levels. If this is intended, adjust max_factor_levels and re-run.")

        if(vartypes[0] == 'numeric'):
            self._problem_type = "regression"
        else:
            self._problem_type = "classification"
            
        n_discard_bits = 0
        if n % 32 != 0:
            n_discard_bits = 32 - n % 32
        if n_discard_bits > 0:
            if n < n_discard_bits:
                raise ValueError("Too few training data points.")
            pad = df.sample(n_discard_bits, axis = 0)
            df = pd.concat([df,pad])
            n = df.shape[0]
            
        df = df.sample(frac=1).reset_index(drop=True)  # shuffle rows
        df.sort_values(by=target_col_name, inplace=True, ignore_index = True)
        indxmat = np.reshape(np.arange(n), (32, int(n/32)))
        indxmat = indxmat.transpose()
        indxmat = np.reshape(indxmat, (n,1))
        df = df.reindex(indxmat.flatten())    

        self._model = brifc.fit(colnames, vartypes, 
                  num_ind, df.iloc[:,num_ind].to_numpy(na_value = self.na_numeric),
                  int_ind, df.iloc[:,int_ind].to_numpy(na_value = self.na_integer),
                  fac_ind, df.iloc[:,fac_ind].to_numpy(na_value = self.na_factor),
                  n, p, self.get_param())
    
    
    def predict(self, test_df, type = 'score'):
        if self._model == None:
            raise ValueError('Model is not available.')
        if not isinstance(test_df, pd.DataFrame):
            raise ValueError("test_df must be a pandas data frame.")
        n = test_df.shape[0]
        p = test_df.shape[1]
        test_colnames = list(test_df.columns.values)
        test_coltypes = test_df.dtypes
        test_vartypes, test_num_ind, test_int_ind, test_fac_ind = _collect_indices(test_coltypes, offset=1)
        brifc.predict(self._model, self.tmp_preddata, test_colnames, test_vartypes,
                             test_num_ind, test_df.iloc[:,[test_num_ind[i] - 1 for i in range(len(test_num_ind))]].to_numpy(na_value = self.na_numeric),
                             test_int_ind, test_df.iloc[:,[test_int_ind[i] - 1 for i in range(len(test_int_ind))]].to_numpy(na_value = self.na_integer),
                             test_fac_ind, test_df.iloc[:,[test_fac_ind[i] - 1 for i in range(len(test_fac_ind))]].to_numpy(na_value = self.na_factor),
                             n, p, self.get_param())

        # read the scores
        scores = pd.read_csv(self.tmp_preddata, sep=' ')
        if(type == 'score'):
            if self._problem_type == 'classification':
                return(scores)
            else:
                head_vec = scores.columns.to_numpy().astype(np.float64)
                pred_mat = scores.to_numpy()
                pred_val = np.dot(pred_mat, head_vec)
                return(list(pred_val))
        else:
            if self._problem_type == 'classification':  
                return(list(scores.idxmax(axis=1)))
            else:
                head_vec = scores.columns.to_numpy().astype(np.float64)
                pred_mat = scores.to_numpy()
                pred_val = np.dot(pred_mat, head_vec)
                return(list(pred_val))                
            
            
    
            
        
        
