
import copy
import warnings
from configargparse import ArgParser
import re
import os
import pandas as pd
import tensorflow as tf
from utils.prefit import make_model
from utils.misc import (read_pickle, sheepish_mkdir,
                        send_message_to_slack, entropy, write_txt)
from utils.constants import SENTENCE_LENGTH, TAGS, STR_VARNAMES
import spacy
from _05_ingest_to_sentence import set_custom_boundaries
import numpy as np
import datetime


class BatchPredictor():
    '''
    Predicts a batch of unlabeled notes, based on PAT_IDs whose last two digits match the digits parameter
    '''
    def __init__(self ,digits, batchstring):
        self.outdir = f"{os.getcwd()}/output/"
        self.datadir = f"{os.getcwd()}/data/"
        self.ALdir = f"{self.outdir}saved_models/AL{batchstring}/"
        self.digits = digits
        self.batchstring = batchstring
        sheepish_mkdir(f"{self.ALdir}final_model/preds")
        
        #################
        # anti-clobber: look for a token or output, write a token if not found
        # if a token or a pred is found, write something to self that short-circuits other methods
        stuff_in_preds_dir = os.listdir(f"{self.ALdir}final_model/preds")
        stuff_with_digits =[i for i in stuff_in_preds_dir if i.split('_')[1][:2] == self.digits]
        if len(stuff_with_digits) > 0:
            self.exit_to_avoid_clobbering = True
            return
        else:
            write_txt("I am a token.  Why must we anthropomorphize everything?",
                      f"{self.ALdir}final_model/preds/token_{self.digits}.txt")
            self.exit_to_avoid_clobbering = False
            #################
            # load data and subset
            self.df = pd.read_csv(f"{self.ALdir}processed_data/full_set/full_df.csv", index_col=0)
            already_pulled_PIDs = []
            for root, dirs, files in os.walk(f"{self.outdir}notes_output"):
                for file in files:
                    already_pulled_PIDs.append(file)
            already_pulled_PIDs = [i for i in already_pulled_PIDs if ('batch' in i or 'AL' in i) and '.txt' in i]
            already_pulled_PIDs = list(set([re.sub(".txt", "", i.split("_")[-1]) for i in already_pulled_PIDs]))

            allnotes = pd.read_pickle(f"{self.outdir}conc_notes_df.pkl")
            allnotes = allnotes.loc[(allnotes.LATEST_TIME < "2019-01-01") &
                                    (~allnotes.PAT_ID.isin(
                                        self.df.PAT_ID)) &  # remove PAT IDs that are already in training set
                                    (~allnotes.PAT_ID.isin(
                                        already_pulled_PIDs))]  # remove PAT IDs that have ever been pulled for a test batch
            allnotes = allnotes.loc[allnotes.PAT_ID.apply(lambda x: x[-2:]) == self.digits]
            allnotes.insert(0, "uid", allnotes.PAT_ID + "_" + allnotes.LATEST_TIME.dt.month.astype(str))
            strdat = pd.read_csv(f"{self.outdir}structured_data_merged_cleaned.csv", index_col=0)
            strdat.insert(0, "uid", strdat.PAT_ID + "_" + strdat.month.astype(str))
            self.allnotes = allnotes.loc[allnotes.uid.isin(strdat.uid)]
            self.strdat = strdat.loc[strdat.uid.isin(self.allnotes.uid)]
            self.str_varnames = [i for i in self.df.columns if re.match("pca[0-9]", i)]
            self.scalers = read_pickle(f"{self.ALdir}/processed_data/full_set/full_sklearn_dict.pkl")
            ####################
            # spacy
            self.sci_nlp = spacy.load("en_core_sci_md", disable=['tagger', 'ner'])
            self.sci_nlp.add_pipe(set_custom_boundaries, before="parser",
                                  name='set_custom_boundaries')
            #######################
            # Load the model and vectorizer
            mod_dict = read_pickle(f"{self.ALdir}/final_model/model_final_{batchstring}.pkl")
            self.model, self.vectorizer = make_model(emb_path=f"{self.datadir}w2v_oa_all_300d.bin",
                                                     sentence_length=SENTENCE_LENGTH,
                                                     meta_shape=len(self.str_varnames),
                                                     tags=TAGS,
                                                     train_sent=self.df['sentence'],
                                                     l1_l2_pen=mod_dict['config']['l1_l2_pen'],
                                                     n_units=mod_dict['config']['n_units'],
                                                     n_dense=mod_dict['config']['n_dense'],
                                                     dropout=mod_dict['config']['dropout'])
            self.model.compile(loss='categorical_crossentropy',
                               optimizer=tf.keras.optimizers.Adam(1e-4))
            self.model.set_weights(mod_dict['weights'])

    def sentencize(self, note, return_sentences=False):
        res = self.sci_nlp(note)
        sentences = []
        s = []
        for i, tok in enumerate(res):
            if (bool(tok.is_sent_start)) & (len(s) > 0):
                newsent = " ".join(s)
                sentences.append(newsent)
                s = []
            s.append(str(tok))
        sentences.append(" ".join(s))
        text = self.vectorizer(np.array([[s] for s in sentences]))
        if return_sentences:
            return text, sentences
        else:
            return text

    def process_structured(self, row, nsent):
        str_mat = pd.concat([row for i in range(nsent)])[STR_VARNAMES]
        str_mat = self.scalers['imputer'].transform(str_mat)
        str_mat = self.scalers['scaler_in'].transform(str_mat)
        str_mat = self.scalers['pca'].transform(str_mat)
        str_mat = self.scalers['scaler_out'].transform(str_mat)
        struc = tf.convert_to_tensor(str_mat, dtype='float32')
        return struc

    def predict(self, ii):
        text = self.sentencize(self.allnotes.combined_notes.iloc[ii])
        struc = self.process_structured(self.strdat.loc[self.strdat.uid == self.allnotes.uid.iloc[ii]],
                                        nsent=text.shape[0])
        pred = self.model.predict([text, struc])
        return pred

    def get_H_stats(self, pred):
        '''pred is a list of predictions for each aspect'''
        nsent = pred[0].shape[0]
        hmat = np.stack([entropy(i) for i in pred])
        average_entropy = np.nanmean(hmat)
        max_entropy = np.nanmax(hmat)
        hflat = hmat.flatten()
        average_entropy_above_median = np.mean(hflat[hflat > np.nanmedian(hflat)])
        hcut = copy.deepcopy(hmat)
        for i in range(hcut.shape[0]):
            hcut[i, :][hcut[i, :] < np.nanmedian(hcut[i, :])] = np.nan
        hcut = np.nanmax(hcut, axis=0)
        average_entropy_above_median_notewise = np.nanmean(hcut)
        return dict(nsent=nsent,
                    average_entropy=average_entropy,
                    max_entropy=max_entropy,
                    average_entropy_above_median=average_entropy_above_median,
                    average_entropy_above_median_notewise=average_entropy_above_median_notewise)

    def loop_predict(self):
        if self.exit_to_avoid_clobbering == True:
            return
        else:
            outlist = []
            for ii in range(len(self.allnotes)):
                pred = self.predict(ii)
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    out = self.get_H_stats(pred)
                out['uid'] = self.allnotes.uid.iloc[ii]
                outlist.append(out)
            dfout = pd.DataFrame(outlist)
            dfout.to_json(f"{self.ALdir}final_model/preds/preds_{self.digits}.json.bz2")
            # os.remove(f"{self.ALdir}final_model/preds/token_{self.digits}.txt") # remove the token, just to keep tidy
            return dfout

    def pull_best_notes(self, n_to_retrieve):
        predfiles = os.listdir(f'{self.ALdir}/final_model/preds/')
        predfiles = [i for i in predfiles if "json.bz2" in i]
        preds = pd.concat([pd.read_json(f"{self.ALdir}/final_model/preds/{i}") for i in predfiles])
        best = preds.loc[(preds.nsent > 20) & (~preds.average_entropy.isna())]
        # deal with multiple PAT IDs:  keep the row with the max entropy
        best = best.reset_index(drop=True)
        best['PAT_ID'] = best.uid.apply(lambda x: x.split("_")[0])
        dups = list(best.PAT_ID.value_counts()[best.PAT_ID.value_counts() > 1].index)
        for d in dups:
            duprows = best.loc[best.PAT_ID == d]
            droppers = duprows.loc[duprows.average_entropy_above_median_notewise != max(
                duprows.average_entropy_above_median_notewise)].index
            best.drop(droppers, inplace=True)
        best = best.sort_values('average_entropy_above_median_notewise').tail(n_to_retrieve)
        assert best.PAT_ID.nunique() == n_to_retrieve
        # merge with other data
        allnotes = pd.read_pickle(f"{self.outdir}conc_notes_df.pkl")
        allnotes = allnotes.loc[(allnotes.LATEST_TIME < "2019-01-01") &
                                (~allnotes.PAT_ID.isin(
                                    self.df.PAT_ID))]  # remove PAT IDs that have ever been pulled for a test batch

        allnotes.insert(0, "uid", allnotes.PAT_ID + "_" + allnotes.LATEST_TIME.dt.month.astype(str))
        strdat = pd.read_csv(f"{self.outdir}structured_data_merged_cleaned.csv", index_col=0)
        strdat.insert(0, "uid", strdat.PAT_ID + "_" + strdat.month.astype(str))
        allnotes = allnotes.loc[allnotes.uid.isin(strdat.uid)]
        strdat = strdat.loc[strdat.uid.isin(allnotes.uid)]
        best = best.merge(allnotes).merge(strdat)

        diags = pd.read_pickle(f"{self.datadir}diagnosis_df.pkl")
        diags = diags.loc[diags.PAT_ID.isin(best.PAT_ID)]
        di = []
        for i in best.PAT_ID.unique():
            d = diags.loc[(diags.PAT_ID == i)].copy()
            d['timediff'] = (best.LATEST_TIME.loc[best.PAT_ID == i].iloc[0] - d['dx_date'])
            qdx = d.loc[(d.timediff <= datetime.timedelta(days=365)) &
                        (d.timediff >= datetime.timedelta(days=0)), 'icd10']

            qdx = ",".join(list(qdx.unique()))
            di.append(dict(PAT_ID=i, qual_dx=qdx))
        tomerge = pd.DataFrame(di)
        best = tomerge.merge(best)
        self.best = best
        best.to_csv(f'{self.ALdir}/final_model/preds/next_batch.csv')
        assert best.shape[0] == n_to_retrieve
        return best

    def write_notes(self):
        sheepish_mkdir(f"{self.outdir}/notes_output/AL{self.batchstring}")
        for i in range(self.best.shape[0]):
            fn = f"AL{self.batchstring}_m{self.best.month.iloc[i]}_{self.best.PAT_ID.iloc[i]}.txt"
            write_txt(self.best.combined_notes.iloc[i],
                      f"{self.outdir}/notes_output/AL{self.batchstring}/{fn}")


def main():
    '''
    Loop through the unlabeled notes and predict.
    '''
    p = ArgParser()
    p.add("-b", "--batchstring", help="the batch number", type=str)

    options = p.parse_args()
    batchstring = options.batchstring
    assert batchstring is not None
    outdir = f"{os.getcwd()}/output/"
    ALdir = f"{outdir}saved_models/AL{batchstring}/"
    sheepish_mkdir(f"{ALdir}final_model/preds")

    digits_vec = np.random.choice([str(i)[1:] for i in range(100, 200)], 100, replace=False)
    for digits in digits_vec:
        try:
            predictor = BatchPredictor(digits=digits, batchstring=batchstring)
            _ = predictor.loop_predict()

        except Exception as e:
            print(e)
            ip = os.popen("/sbin/ip -o -4 addr list eth0 | awk '{print $4}' | cut -d/ -f1").read()
            send_message_to_slack(f"problem processing {digits} at {ip}")


if __name__ == "__main__":
    main()

