import logging
import re
from typing import Optional

import numpy as np
import pandas as pd

from icdpicpy.src.datamodel import Constants, Table, Score
from icdpicpy.src.datamodel.CodeMappingMethod import ICD10CodeMappingMethod, ISSCalcMethod
from icdpicpy.src.repository.DataRepository import DataRepository


class CalcScore:
    def __init__(self, data_repo: DataRepository):
        self.__logger = logging.getLogger(self.__class__.__name__)
        self.__data_repo = data_repo

    def execute(self, patient_code_dict: dict, score_list: list = None, use_icd10: bool = True,
                icd10_method: str = 'gem_max', iss_calc_method: int = 1) -> Optional[pd.DataFrame]:
        self.__logger.debug(
            f'execute for {len(patient_code_dict)} patients, scores {score_list}, use_icd10: {use_icd10},'
            f' icd10 method: {icd10_method}, iss calculation method:{iss_calc_method}')

        if not self.__validate_params(patient_code_dict, use_icd10, icd10_method, iss_calc_method, score_list):
            return
        if score_list is None:
            score_list = Score.all_scores
        iss_calc_method = ISSCalcMethod.from_int(iss_calc_method)
        severity_data_df = self.__data_repo.read_data_file(Table.SeverityMap.file_name)
        if use_icd10:
            severity_data_df = self.__calc_icd10_severity(patient_code_dict, severity_data_df, icd10_method)
        res_df = None
        if severity_data_df is not None and Score.ISS_score in score_list:
            patient_iss_df = self.__calc_iss(patient_code_dict, severity_data_df, iss_calc_method)
            res_df = patient_iss_df
        if severity_data_df  is not None and Score.NISS_score in score_list:
            patient_niss_df = self.__calc_niss(patient_code_dict, severity_data_df, iss_calc_method)
            res_df = patient_niss_df if res_df is None else \
                res_df.merge(patient_niss_df, on=Table.Result.patient)
        if Score.Injury_cause in score_list:
            injury_cause_df = self.__data_repo.read_data_file(Table.InjuryCause.i9_file_name, sep='\t')
            if use_icd10:
                injury_cause_icd10_df = self.__data_repo.read_data_file(Table.InjuryCause.i10_file_name, sep='\t')
                injury_cause_df = pd.concat([injury_cause_df, injury_cause_icd10_df])

            patient_injury_cause_df = self.__get_injury_cause(patient_code_dict, injury_cause_df)
            if patient_injury_cause_df is not None:
                res_df = patient_injury_cause_df if res_df is None else \
                    res_df.merge(patient_injury_cause_df, on=Table.Result.patient)
        return res_df

    def __calc_iss(self, patient_code_dict: dict, severity_df: pd.DataFrame,
                   iss_calc_method: ISSCalcMethod) -> pd.DataFrame:
        self.__logger.debug('calc ISS')
        res = {Table.Result.patient: [], Table.Result.iss: []}
        for patient, codes in patient_code_dict.items():
            codes = [self.__prepare_code(c) for c in codes]
            # get only codes of the current patient
            curr_severity_df = severity_df[severity_df[Table.SeverityMap.code].isin(codes)]
            # severity score of 9 implies unknown severity.
            # Thus we want to exclude these as long as there is at least one known severity for the body region
            # However if all severity scores for the body region are 9 then we will assign max value of 9
            for br in curr_severity_df[Table.BaseColumns.issbr].unique():
                unique_values = curr_severity_df[Table.BaseColumns.severity].unique()
                if len(unique_values) > 1 and Constants.UnknownSeverityValue in unique_values:
                    # there are other values so Unknown can be removed
                    known_values_cond = (curr_severity_df[Table.BaseColumns.issbr] != br) | \
                                        (curr_severity_df[Table.BaseColumns.severity] != Constants.UnknownSeverityValue)
                    curr_severity_df = curr_severity_df.loc[known_values_cond, :]
            # find max severity for each body region
            br_df = curr_severity_df.groupby([Table.BaseColumns.issbr])[Table.BaseColumns.severity].max()
            # find iss as a sum of square of the largest severity scores but 75 is max
            severity_values: list = br_df.reset_index()[Table.BaseColumns.severity].tolist()
            severity_values.sort(reverse=True)
            max_pos = min(3, len(severity_values))
            severity_values = severity_values[:max_pos]
            if all([x == Constants.UnknownSeverityValue for x in severity_values]):
                iss_value = np.NaN
            else:
                severity_values = [0 if x == Constants.UnknownSeverityValue else x
                                   for x in severity_values]
                sqr_sum = 0
                if Constants.SeverityMaxValue in severity_values:
                    if iss_calc_method == ISSCalcMethod.default:
                        sqr_sum = Constants.MaxISSValue
                    elif iss_calc_method == ISSCalcMethod.extreme_cut:
                        severity_values = [x - 1 if x == Constants.MaxISSValue else x
                                           for x in severity_values]
                        sqr_sum = sum([x * x for x in severity_values])
                else:
                    sqr_sum = sum([x * x for x in severity_values])
                iss_value = min(sqr_sum, Constants.MaxISSValue)
            res[Table.Result.patient].append(patient)
            res[Table.Result.iss].append(iss_value)
        res_df = pd.DataFrame.from_dict(res)
        return res_df

    def __calc_niss(self, patient_code_dict: dict, severity_df: pd.DataFrame,
                    iss_calc_method: ISSCalcMethod) -> pd.DataFrame:
        self.__logger.debug('calc NISS')
        res = {Table.Result.patient: [], Table.Result.niss: []}
        for patient, codes in patient_code_dict.items():
            codes = [self.__prepare_code(c) for c in codes]
            # get only codes of the current patient
            severity_values = severity_df.loc[severity_df[Table.SeverityMap.code].isin(codes),
                                              Table.SeverityMap.severity].tolist()
            # severity score of 9 implies unknown severity.
            # Thus we want to exclude these as long as there is at least one known severity for the patient
            # However if all severity scores for the patient are 9 then we will assign value NA
            if all([x == Constants.UnknownSeverityValue for x in severity_values]):
                niss_value = np.NaN
            else:
                severity_values = [0 if x == Constants.UnknownSeverityValue else x
                                   for x in severity_values]
                severity_values.sort(reverse=True)
                max_pos = min(3, len(severity_values))
                severity_values = severity_values[:max_pos]
                if iss_calc_method == ISSCalcMethod.default:
                    if any([x == Constants.SeverityMaxValue for x in severity_values]):
                        niss_value = Constants.MaxISSValue
                    else:
                        niss_value = sum([x * x for x in severity_values])
                elif iss_calc_method == ISSCalcMethod.extreme_cut:
                    severity_values = [x - 1 if x == Constants.SeverityMaxValue else x
                                       for x in severity_values]
                    niss_value = sum([x * x for x in severity_values])
                else:
                    niss_value = 0
                niss_value = min(niss_value, Constants.MaxISSValue)

            res[Table.Result.patient].append(patient)
            res[Table.Result.niss].append(niss_value)

        res_df = pd.DataFrame.from_dict(res)
        return res_df

    def __prepare_code(self, code: str) -> str:
        code = code.strip().replace(".", "")
        if len(code) > 5 and code[5] == "X":
            return code[:5]
        if len(code) > 6 and code[6] == "X":
            return code[:6]
        return code

    def __validate_params(self, patient_code_dict: dict, use_icd10: bool, icd10_method: str,
                          iss_calc_method: int, score_list: Optional[list]) -> bool:
        self.__logger.debug('validate parameters')
        if not patient_code_dict:
            self.__logger.error(f'No patient data')
            return False

        if use_icd10 and ICD10CodeMappingMethod.from_string(icd10_method) is None:
            self.__logger.error(f'Unknown ICD10 mapping method {icd10_method}. '
                                f'Value should be one of {ICD10CodeMappingMethod.all_values()}')
            return False
        if ISSCalcMethod.from_int(iss_calc_method) is None:
            self.__logger.error(f'Unknown ISS calculation method {iss_calc_method}. '
                                f'Value should be one of {ISSCalcMethod.all_values()}')
            return False
        if score_list is not None:
            if any([x not in Score.all_scores for x in score_list]):
                self.__logger.error(f'Unknown score: {[x for x in score_list if x not in Score.all_scores]}. '
                                    f'Score can be one of {Score.all_scores}')
                return False
        return True

    def __get_injury_cause(self, patient_code_dict: dict, injury_cause_df: pd.DataFrame) -> Optional[pd.DataFrame]:
        self.__logger.debug('get injury cause')
        injury_cause_df[Table.InjuryCause.code] = injury_cause_df[Table.InjuryCause.code].apply(lambda x: x.strip())

        res = {Table.Result.patient: list(patient_code_dict.keys()),
               Table.Result.mechanism: [], Table.Result.intention: []}
        for p, codes in patient_code_dict.items():
            codes = [self.__prepare_code(c) for c in codes]
            mechanism, intention = self.__get_mechanism_and_intention(codes, injury_cause_df)
            res[Table.Result.mechanism].append(mechanism)
            res[Table.Result.intention].append(intention)
        return pd.DataFrame(data=res)

    def __get_mechanism_and_intention(self, codes: list, injury_cause_df: pd.DataFrame) -> tuple:
        self.__logger.debug(f'get mechanism and intention for {codes}')
        res = []
        for c in codes:
            # match by first 3 character to cut the most variants
            curr_df = injury_cause_df[injury_cause_df[Table.InjuryCause.code].str[:3] == c[:3]]
            # substitute 'X' in the end of a code on '*' for regex
            curr_df[Table.InjuryCause.code] = curr_df[Table.InjuryCause.code].apply(
                lambda x: f'{x[:-1]}.*' if x[-1] == 'X' else x
            )
            # substitute 'X' in the middle of a code on '(.)' for regex
            curr_df[Table.InjuryCause.code] = curr_df[Table.InjuryCause.code].apply(
                lambda x: f'{x[:x.rfind("X")]}(.){x[x.rfind("X") + 1:]}' if x.rfind('X') > 0 else x
            )
            # match codes
            match = curr_df[Table.InjuryCause.code].apply(lambda x: re.search(x, c) is not None)
            curr_df = curr_df[match]
            if curr_df.empty:
                continue
            curr_df[Table.InjuryCause.code] = c
            curr_df[Table.InjuryCause.mechanism] = curr_df. \
                groupby([Table.InjuryCause.code, Table.InjuryCause.intention])[Table.InjuryCause.mechanism]. \
                transform(lambda x: '/'.join(x))
            curr_df = curr_df.drop_duplicates()
            res.append(curr_df)
        mechanism = ''
        intention = ''
        if res:
            df = pd.concat(res)
            if not df.empty:
                mechanism = '/'.join(df[Table.InjuryCause.mechanism].unique().tolist())
                intention = '/'.join(df[Table.InjuryCause.intention].unique().tolist())
        return mechanism, intention

    def __calc_icd10_severity(self, patient_code_dict: dict, severity_data_df: pd.DataFrame,
                              icd10_method: str) -> Optional[pd.DataFrame]:
        self.__logger.debug('calc icd10 severity')
        icd10_icd9_map_df = self.__data_repo.read_data_file(Table.Icd10Icd9Map.file_name)
        all_codes = [item for _, values in patient_code_dict.items() for item in values]
        all_codes = [self.__prepare_code(x) for x in all_codes]
        icd10_icd9_map_df = icd10_icd9_map_df[icd10_icd9_map_df[Table.Icd10Icd9Map.icd10].isin(all_codes)]
        # map icd10 on severity
        merged_severity_data_df = severity_data_df.merge(icd10_icd9_map_df,
                                                         left_on=Table.SeverityMap.code,
                                                         right_on=Table.Icd10Icd9Map.icd9)
        i10_severity_df = merged_severity_data_df[[Table.Icd10Icd9Map.icd10, Table.SeverityMap.severity,
                                                   Table.SeverityMap.issbr]]
        i10_severity_df = i10_severity_df.rename({Table.Icd10Icd9Map.icd10: Table.SeverityMap.code}, axis=1)
        # concat icd 9 and icd10 severity tables
        severity_data_df = pd.concat([severity_data_df, i10_severity_df])
        severity_data_df = severity_data_df[severity_data_df[Table.SeverityMap.code].isin(all_codes)]
        # resolve mapping conflict according to icd10_method
        method = ICD10CodeMappingMethod.from_string(icd10_method)
        if method is None:
            self.__logger.error(f'Unknown ICD10 mapping method {icd10_method}')
        if method == ICD10CodeMappingMethod.gem_min:
            agg_function = 'min'
        elif method == ICD10CodeMappingMethod.gem_max:
            agg_function = 'max'
        else:
            self.__logger.error(f'Unknown icd10 mapping method {icd10_method}')
            return None
        severity_data_df[Table.SeverityMap.severity] = \
            severity_data_df.groupby([Table.SeverityMap.code])[Table.SeverityMap.severity].transform(agg_function)
        severity_data_df = severity_data_df.drop_duplicates()
        return severity_data_df
