# -*- coding: utf-8 -*-
"""
Created on Sun Oct 14 23:31:49 2018

@author: hmelberg
"""

import copy
import pickle
import pandas as pd

from rulebook.utils import _listify, _expand_rulelist, OneRule

from rulebook.check import evaluate


# %%
class RuleGroup():
   def __init__(self, rules, cols=None, name=None,
                description=None, action=None, comments=None, etc=None):
       self.rules = _listify(rules)
       self.cols = _listify(cols)
       self.name = name
       self.description = description
       self.action = action
       self.comments = comments
       self.etc = etc


# %%
class RuleBook():
   def __init__(self, rules=None):
       self.rulebook_list = list()
       self.comments = None,
       self.info = None,

   def add(self, rules, cols=None, name=None, description=None,
           action=None, comments=None, etc=None):
       """
       Adds a rule object (one or more rules) to the rulebook

       args:
           - rule (str, list of str): The name of the function determining valid observations
                   - can use inbuilt functions of self-defined functions
                   - can add arguments to functions in parenthesis, but must use keywords
                   - can use a list of one or more functions and columns: one or multiple rules can apply to one or more columns
                   - the function should return a series with 1 for valid rows and 0 for invalid rows (missing = invalid = false = 0)
           - cols (str, list of str): The column(s) where the rule(s) apply
                   - Single columns can be a str, and a comma seperated str is converted to a list with mupliple cols
           - name (str, optional): The rule can be assigned a name. if no name is assigned, one will be autogenerated (rule_1 etc)
           - description (str, optional): Describe rule in more detail
           - action (str, list of str): The name  of the function to be used  if rule is violated  (arguments in allowed in parenthesis)

       examples
           - add('is_int', 'age')
           - add(['is_int', 'is_positive'], ['age', 'days'])
           - add('length(max=5)', 'icd')

       """

       rule_group = RuleGroup(rules=rules, cols=cols, name=name, description=description,
                              action=action, comments=comments, etc=etc)

       # give the rule a name if no name is given
       if not rule_group.name:
           nrules = len(self.rulebook_list)
           proposed_name = 'rule_' + str(nrules + 1)
           existing_names = [rule.name for rule in self.rulebook_list]

           if proposed_name in existing_names:
               existing_nums = [int(name.split('_')[1]) for name in existing_names if
                                'rule_' in name]  # potentialproblem a name such as rule_no_negative
               max_num = max(existing_nums)
               proposed_name = 'rule_' + str(max_num + 1)
           rule_group.name = proposed_name

       # convert col to list if it is a string with , inside
       if rule_group.cols:
           if ',' in rule_group.cols:
               cols = rule.cols.split(',')
               cols = [strip(col) for col in cols]
               rule_group.cols = cols

       setattr(self, rule_group.name, rule_group)
       self.rulebook_list.append(rule_group)

   def add_many(self, rules):
       """
       add many rules to the rule object

       Note
           - Only works when column name is part of the expressions (ie. not with functions since functions also need to be assigned to columns)
           - Both a list and a comma separated string is ok

           add_many('age>18, birth_year>1970')
           add_many(['age>18', 'birth_year>1970'])
       """

       if isinstance(rules, str):
           rules = rules.split(',')
       for rule in rules:
           self.add(rule=rule)

   def delete(self, rules):
       if isinstance(rules, str): rules = set([rules])
       keep_rules = [rule for rule in self.rules if rule.name not in rules]
       self.rules = keep_rules

   def append_col(self, name, cols):
       cols = listify_cols
       modified_rules = []
       for rule in self.rulebook_list:  # rulebook_lilst should be private  .. and called sth else
           if rule.name == name:
               rule.cols.extend(cols)
           modified_rules.append(rule)
       return modified_rules

   def drop_col(self, name, cols):
       cols = listify_cols
       modified_rules = []
       for rule in self.rulebook_list:  # rulebook_lilst should be private  .. and called sth else
           if rule.name == name:
               existing_cols = rule.cols
               keep_cols = existing_cols - set(cols)
               rule.cols = keep_cols
           modified_rules.append(rule)
       return modified_rules

   def view(self, rules=None):
       if not rules: rules = self.rulebook_list
       for rule in rules:
           print(rule.name, rule.rules)

   def view_cols(self, cols=None, rules=None):
       self._fix_format()

       if not rules: rules = self.rulebook_list

       col_rules = {}
       for rule in rules:
           for col in rule.cols:
               if (not cols) or (cols in rule.cols):
                   if col in col_rules:
                       col_rules[col].append(rule)
                   else:
                       col_rules[col] = [rule]

       for col, rules in col_rules.items():
           for rule in rules:
               print(col, rule.name, rule.rule)

   def save(self, file):
       with open(file, 'wb') as output:
           pickle.dump(self, output, protocol=0)
           # pickle.HIGHEST_PROTOCOL

   def _check(self, df, rules=None, cols=None, change=False, out='report'):

       # use all or only some rules
       all_rules = set(self.rulebook_list)
       all_names = set([rule.name for rule in all_rules])

       if rules:
           specified_names = set(_listify(rules))
           subset_names = all_names.intersection(specified_names)
       else:
           subset_names = all_names

       subset_rules = [rule for rule in all_rules if rule.name in subset_names]

       # use all or only some columns
       if cols:
           subset_rules = copy.deepcopy(subset_rules)
           include_cols = set(cols)

           for rule in subset_rules:
               existing_cols = set(rule.cols)
               new_cols = existing_cols.intersection(include_cols)
               rule.cols = list(new_cols)

       # same rule(s) may apply to many column(s), so make
       # separate rule for each combination
       expanded_rules = _expand_rulelist(subset_rules)

       # dictionary to save info from check
       result = {}

       for rule in expanded_rules:
           ok = globals()[rule.func[0]](df=df, **rule.args[0])  # hmmm why are these tuples?)

           if rule.args[0]['col']:
               col = rule.args[0]['col']
               name = f"{rule.name}_col{col}"

               result[name] = dict()
               result[name]['ok'] = ok
               result[name]['nans'] = df[col].isnull().sum()
               result[name]['ninvalid'] = (~ok).sum()
               result[name]['values'] = df[col][~ok].values
               result[name]['series'] = df[col][~ok]
               result[name]['df'] = df[~ok]
           else:
               result[rule.name] = dict()

               result[name]['ok'] = ok
               result[name]['nans'] = ok.isnull().sum()
               result[name]['ninvalid'] = (~ok).sum()
               result[name]['df'] = df[~ok]

           # change the column(s)
           if change and rule.action:
               try:
                   new_df = globals()[rule.action[0]](df=df, ok=ok, **rule.action_args[0])
               except:
                   print(f"Error for rule {rule.name, rule.text}")

       if out == 'report':
           for name, res in result.items():
               print(name, res['ninvalid'], res['nans'])
           return
           # return fails

       elif out == 'results':
           for name, res in result.items():
               print(name, res['ninvalid'], res['nans'])
           return result

       elif out == 'change':
           return new_df

       elif out == 'check_and_change':
           return result, df

       else:
           print('out argument is wrong')
       return

   def check(self, df, rules=None, cols=None):
       results_out = self._check(df=df, rules=rules, cols=cols, out='report')
       return results_out

   def change(self, df, rules=None, cols=None):
       changed_df = self._check(df=df, rules=rules, cols=cols, out='change', change=True)
       return changed_df

   def check_change_check(self, df, rules=None, cols=None):
       df_with_fails_before = self._check(df=df, rules=rules, cols=cols)
       changed_df = self._check(df=df, rules=rules, cols=cols, out='change')
       df_with_fails_after = self._check(df=changed_df, rules=rules, cols=cols)
       return changed_df, df_with_fails_before, df_with_fails_after


def load_rulebook(file):
   with open(file, 'rb') as input:
       rb = pickle.load(input)
   return rb


def suggest_rules(df, rules=None, cols=None, sample=0.1, pid='pid', threshold=0.1):
   """
   Suggest a rule book when almost all of the results conform to a rule
   """
   rule_col = {}
   # df = df.sample(sample)

   all_cols = list(df.columns)
   # rb = RuleBook()

   # first find dtypes
   dtypes = df.dtypes.apply(lambda x: x.name).to_dict()

   num_cols = list(df.select_dtypes(include='number').columns)
   obj_cols = list(df.select_dtypes(include='object').columns)
   date_cols = list(df.select_dtypes(include='datetime').columns)
   cat_cols = list(df.select_dtypes(include='category').columns)
   int_cols = list(df.select_dtypes(include='int').columns)
   float_cols = list(df.select_dtypes(include='float').columns)
   bool_cols = list(df.select_dtypes(include='bool').columns)
   # delta_cols
   # other_cols
   # mixed_cols

   for col in all_cols:
       rule_col[col] = [f'dtype_{dtypes[col]}']

   # find datecols
   # df.infer_objects().dtypes
   dtype = {col: [] for col in all_cols}
   obs = df.head(1000)

   obj_cat_cols = obj_cols + cat_cols

   for col in obj_cat_cols:

       ser = obs[col].dropna()

       if ser.empty or all(ser == ' '):
           continue

       print(col)
       ok = pd.to_numeric(ser, errors='coerce')
       if ok.notnull().sum() > 500:
           rule_col[col].append('dtype_num')

       ok = pd.to_datetime(ser.dropna().head(1000), errors='coerce')
       if ok.notnull().sum() > 500:
           dtype[col] = 'date'
           rule_col[col].append('dtype_date')

       uniques = ser.nunique()
       #print(col, uniques)
       if uniques:
           if len(obs) / uniques > 3:
               dtype[col] = 'category'
               rule_col[col].append('dtype_category')  # and delete if existing is rule that says object

       if pid:
           ok = (obs.groupby(pid)[col].nunique() == 1)
           if ok.sum() > 0.95 * obs[pid].nunique():
               rule_col[col].append('stable')

       if ser.nunique() > 990:
           rule_col[col].append('unique')

       len_dist = ser.str.len().value_counts()
       # outliers out, frequency of ? relative to n, also dispersion in general? need measure here, statistics

       if ser.nunique() > 990:
           rule_col[col].append('unique')

   for col in num_cols:
       #print(col)
       ser = obs[col].dropna()

       if ser.empty:
           continue

       if ser.nunique() < 3:
           dtype[col] = 'bool'
           rule_col[col].append('bool')

       if ser.min() > 0 and ser.max() < 1:
           rule_col[col].append('between(0,1)')

       if pid:
           ok = (obs.groupby(pid)[col].nunique() == 1)

           if ok.sum() > 0.95 * obs[pid].nunique():
               rule_col[col].append('stable')  # (maybe unique(groupby='pid'))

       if col.startswith('age'):
           rule_col[col].extend(['always_positive', "increase(groupby='pid', sort='date'), max(150)"])

       if obs[col].nunique() > 995:
           rule_col[col].append('unique')

       if col == 'pid':
           rule_col[col].append('no_missing')

       if ser.sort_values().iloc[5] >= 0:
           rule_col[col].append('is_positive')

   # add special rules - examine dtypes, max, min, contains etc

   return rule_col