"""
Medical lexicon NLP extraction pipeline

File contains: Compares the validation set with the NLP pipeline's labeling and outputs some relevant statistics afterwards in a CSV-style table.

-- (c) Rene Ahlsdorf 2019 - Team  D in the HST 953 class
"""

from na_pipeline_tool.utils import logger
from na_pipeline_tool.utils import config
from na_pipeline_tool.utils import helper_classes
from na_pipeline_tool import utils 

from na_pipeline_tool.utils import progressbar

import re
import pandas as pd
from joblib import Parallel, delayed  

import sys, os
import collections
import numpy as np 

import sklearn.metrics

class ValidationTableModule(helper_classes.Module):

    def __init__(self):
        super().__init__()

        self._validation_set = config.get_pipeline_config_item(self.module_name(), 'validation_set_file', None)
        self._df_notes_labeled_paths = config.get_pipeline_config_item(self.module_name(), 'input_note_files', [])

        self._loaded_df = []
        self._compare_df = None
        self._orig_df = None
        self._loaded_validation = None
        self._loaded_validation_labels = None
        self._loaded_validation_label_map = None

        logger.log_info('Loading validation note labeling file')
        self._loading_validation_labeling_file()
        logger.log_info('DONE: Loading validation note labeling file')

        logger.log_info('Loading NLP pipeline processed note files')
        self._loading_note_files()
        logger.log_info('DONE: NLP pipeline processed note files')

        logger.log_info('Computing and outputting statistics')
        
        line_list=[]
        for _ in self._loaded_df:
            line_list.append(';')
            table = self._do_statistics(_)
            for _r in range(len(table[0])):
                elems = [_c[_r] for _c in table]
                line_list.append(';'.join(elems))
            line_list.append(';')
            line_list.append(';')
        
        logger.log_info('CSV Table Output:')

        for _ in line_list:
            print(_)

    def _loading_note_files(self):
        if not self._df_notes_labeled_paths:
            raise RuntimeError('Please specify valid note input files.')
        
        def load_file(path):
            filename = utils.default_dataframe_name(path)
            assert os.path.isfile(filename), 'Could not find note parquet file: {}'.format(filename)
            df = pd.read_parquet(filename)
            df.columns = [_.upper() for _ in df.columns]

            assert 'ROW_ID' in list(df.columns), 'Notes file need to have columns: Row_id, predicted_categories'
            assert 'PREDICTED_CATEGORIES' in list(df.columns), "Processed note file needs to have the PREDICTED_CATEGORIES column generated by e.g. the negation module."
            df['PREDICTED_CATEGORIES'] = df.PREDICTED_CATEGORIES.str.upper()
            df['PREDICTED_CATEGORIES'] = df.PREDICTED_CATEGORIES.str.replace(' ', '_')
            df['PREDICTED_CATEGORIES'] = df.PREDICTED_CATEGORIES.str.split('|')
            if 'FOUND_EVIDENCE' in list(df.columns):
                df['FOUND_EVIDENCE'] = df['FOUND_EVIDENCE'].astype(bool)
                df = df[df['FOUND_EVIDENCE']]
            
            return df

        for _ in self._df_notes_labeled_paths:
            self._loaded_df.append(load_file(_))

        unique_labels = []
        all_pred_cats = []
        for _ in self._loaded_df:
            all_pred_cats.extend(list(_.PREDICTED_CATEGORIES))

        for _ in [*all_pred_cats, self._loaded_validation_labels]:
            unique_labels.extend(_)

        unique_labels = set(unique_labels)
        
        lbl_id = 3
        self._loaded_validation_label_map = {"NONE" : 1, "Any" : 2}
        for _lbl in unique_labels:
            if _lbl == "NONE":
                continue
            if _lbl == "Any":
                continue
            self._loaded_validation_label_map[_lbl] = lbl_id
            lbl_id += 1

        for _ in self._loaded_df:
            _['PREDICTED_CATEGORIES'] = _.PREDICTED_CATEGORIES.apply(lambda x: [self._loaded_validation_label_map[_] for _ in x])
            _['PREDICTED_CATEGORIES'] = _.PREDICTED_CATEGORIES.apply(lambda x: [*x, 2] if not 1 in x else x)

        self._loaded_validation['NOTE_TYPES'] = self._loaded_validation.NOTE_TYPES.apply(lambda x: [self._loaded_validation_label_map[_] for _ in x])
        self._loaded_validation['NOTE_TYPES'] = self._loaded_validation.NOTE_TYPES.apply(lambda x: [*x, 2] if not 1 in x else x)

    
    def _loading_validation_labeling_file(self):
        assert self._validation_set, 'Please specify a validation labeling file.'
        try:
            with open(self._validation_set, 'r') as file:
                self._loaded_validation = file.readlines()
                self._loaded_validation = self._loaded_validation[1:]
                self._loaded_validation = [_.strip() for _ in self._loaded_validation]
                self._loaded_validation = [_.split(',') for _ in self._loaded_validation]
                self._loaded_validation = [[int(_[0]), [_.upper().replace(' ', '_') for _ in str(_[1]).split('|')], (int(_[2]) > 0)] for _ in self._loaded_validation]
                self._loaded_validation = pd.DataFrame(self._loaded_validation, columns=['ROW_ID', 'NOTE_TYPES', 'VALID_INCLUDED'])
                self._loaded_validation.loc[~self._loaded_validation['VALID_INCLUDED'], 'NOTE_TYPES'] = pd.Series([['NONE']]*self._loaded_validation.shape[0])
        except:
            raise RuntimeError('Error while processing validation labeling file. Check file structure.')
        self._loaded_validation_labels = []
        for _i, _ in self._loaded_validation.iterrows():
            self._loaded_validation_labels.extend(_['NOTE_TYPES'])

        self._loaded_validation_labels = set(self._loaded_validation_labels)


    def dump_examples_for_comparison(self):
        if not self._compare_df:
            logger.log_warn('Could not find comparison df - Skipping dumping of exemplary notes.')
            return
        
        self._get_examples_for_categories = [_.upper() for _ in self.get_examples_for_categories]
        if not self._get_examples_for_categories:
            logger.log_warn('No categories specified for dumping exemplary sentences.')
            return
        
        unknown_categories = [_ for _ in self._get_examples_for_categories if not _ in [*self._get_examples_for_categories, 'NO_FINDING']]
        if unknown_categories:
            logger.log_warn('The following categories are not present in the provided dataframes: {}'.format(unknown_categories))
            return
        
        example_list = []
        # for _cat in self._get_examples_for_categories:
        #     # Get example sentences

    def _do_statistics(self, df_cmp):
        validset = self._loaded_validation.sort_values('ROW_ID').reset_index(drop=True)[['ROW_ID', 'NOTE_TYPES']].copy()
        validset = validset.drop_duplicates(subset=['ROW_ID'])

        predicted = df_cmp[['ROW_ID', 'PREDICTED_CATEGORIES']].copy()
        predicted = predicted.rename(columns={'PREDICTED_CATEGORIES' : 'PREDICTED_CAT'})
        predicted = predicted.drop_duplicates(subset=['ROW_ID'])

        validset = validset.merge(predicted, how='left', on='ROW_ID')
        validset.loc[validset['PREDICTED_CAT'].isnull(), 'PREDICTED_CAT'] = pd.Series([[1]]*validset.shape[0])
        validset.loc[validset['NOTE_TYPES'].isnull(), 'NOTE_TYPES'] = pd.Series([[1]]*validset.shape[0])
        
        validset['MATCHED'] = validset.apply(lambda x: [_ for _ in x.NOTE_TYPES if _ in x.PREDICTED_CAT], axis=1)
        validset['UNMATCHED_VALID'] = validset.apply(lambda x: [_ for _ in x.NOTE_TYPES if _ not in x.PREDICTED_CAT], axis=1)
        validset['UNMATCHED_PREDICTED']= validset.apply(lambda x: [_ for _ in x.PREDICTED_CAT if _ not in x.NOTE_TYPES], axis=1)
        validset['CORRECT_NOTE'] = False
        #validset['CORRECT_NOTE_ANY'] = True

        validset.loc[(validset.UNMATCHED_VALID.str.len() == 0) & (validset.UNMATCHED_PREDICTED.str.len() == 0), 'CORRECT_NOTE'] = True
        #validset.loc[((validset.UNMATCHED_VALID.str.len() == 0) & (validset.UNMATCHED_PREDICTED.str.len() == 0), 'CORRECT_NOTE_ANY'] = True
       
        max_index = max(self._loaded_validation_label_map.values())

        one_hot_valid = np.zeros((validset.shape[0], max_index))
        one_hot_pred  = np.zeros((validset.shape[0], max_index))

        _i = 0
        for _, _row in validset.iterrows():
            predicted = [_-1 for _ in _row.PREDICTED_CAT if _ != 0]
            valid = [_-1 for _ in _row.NOTE_TYPES if _ != 0]

            if valid:
                one_hot_valid[_i, valid] = 1
            if predicted:
                one_hot_pred[_i, predicted] = 1

            _i += 1
            

        assert len(validset.groupby('ROW_ID').first().reset_index()) == len(validset)

        ##logger.log_info('Correctly identified notes: {}/{} ({}%)'.format(validset['CORRECT_NOTE'].sum(), len(validset), validset['CORRECT_NOTE'].sum() * 100 / len(validset)))


        def print_report(y_true, y_pred):
            return sklearn.metrics.precision_score(y_true, y_pred), sklearn.metrics.recall_score(y_true, y_pred), sklearn.metrics.accuracy_score(y_true, y_pred), sklearn.metrics.f1_score(y_true, y_pred, average='micro'), sklearn.metrics.f1_score(y_true, y_pred, average='macro'), sklearn.metrics.f1_score(y_true, y_pred, average='binary') #, np.sum(y_true == y_pred),'/', len(y_pred), '({}, {})'.format(y_true.sum(), y_pred.sum())

        cats = sorted(self._loaded_validation_label_map.keys())

        table = [['Lexicon', 'Precision', 'Recall', 'Accuracy', 'F1 Micro', 'F1 Macro', 'F1 Binary']] + [[_] for _ in cats] + [[]]
        
        for _col, _cat in enumerate(cats):#range(1, max_index + 1):
            label = [_[1] for _ in self._loaded_validation_label_map.items() if _cat == _[0]][0]
            _p, _r, _a, _f1_mi, _f1_ma, _f1_bin = print_report(one_hot_valid[:, label - 1],one_hot_pred[:, label - 1])
            table[_col + 1].extend(['{:.2f}%'.format(_p*100),'{:.2f}%'.format(_r*100),'{:.2f}%'.format(_a*100), '{:.4f}'.format(_f1_mi), '{:.4f}'.format(_f1_ma), '{:.4f}'.format(_f1_bin)])
        
        table[_col+2] = ['Total','','','{:.2f}%'.format((validset['CORRECT_NOTE'].sum() * 100 / len(validset))), '','', ''] #'{:.2f}%'.format((validset['CORRECT_NOTE_ANY'].sum() * 100 / len(validset)))
        return table
  
    @classmethod
    def register_argparser_object(cls, subparser_instance):
        subparser_instance.add_parser(cls.module_name(), help="Compares manually labeled notes with the framework's labeling and outputs all relevant statistics in a CSV table style.")
    
