from sklearn.base import TransformerMixin, ClassifierMixin

from value_extractor_transformer import EFTransformer, QRSTransformer, LBBBTransformer, SinusRhythmTransformer, NYHATransformer

class ClinicalDecisionModel(TransformerMixin, ClassifierMixin):

    """

    Replicate prediction of clinical descion tree

    in 'Circulation-2012-Tracy-1784-800.pdf' Appendix 3

    Does not need to be fit and transforms data on its own

    """

    LVEF = 'lvef'

    QRS = 'qrs'

    LBBB = 'lbbb'

    SINUS_RHYTHM = 'sr'

    NYHA = 'nyha'

    def __init__(self):

        self.transformers = dict()

        self.transformers[self.LVEF]= EFTransformer('max',3)

        self.transformers[self.QRS]= QRSTransformer('mean', 3)

        self.transformers[self.LBBB]= LBBBTransformer()

        self.transformers[self.SINUS_RHYTHM]= SinusRhythmTransformer()

        self.transformers[self.NYHA]= NYHATransformer()

    def fit(self, X = None, y = None):

        for trans_key in self.transformers:

            self.transformers[trans_key].fit(X, y)

        return self

    def transform(self, X):

        return [self.__find_values(empi) for empi in X]

    def __find_values(self, empi):

        values = dict()

        for trans_key in self.transformers:

            values[trans_key] = self.transformers[trans_key].get_feature(empi)

        return values

    def predict(self, X):

        X_transformed = self.transform(X)

        predicted_colors = map(self.predict_color, X_transformed)

        y_hat = map(self.prediction_from_color, predicted_colors)

        return y_hat

    def predict_color(self, values):

        """

        Summary of logic (all have ef > 35):

        NYHA I:

            Orange - ef < 30, qrs > 150, lbbb

            Red -    ef > 30 or qrs < 150 or no lbbb

        NYHA II:

            Green -  qrs > 150, lbbb, sinus rhythm

            Yellow - qrs < 150, lbbb, sinus rhythm

            Orange - qrs > 150, no lbbb

            Red -    no lbbb and qrs < 150 or lbbb and no SR

        NYHA III:

            Green -  qrs > 150, lbbb, sinus rhythm

            Yellow - qrs > 150, no lbbb, sinus rhythm

            Orange - qrs < 150, no lbbb, sinus rhythm

            Red -    qrs < 120 or no sinus rhythm

        NYHA IV:

            Red -    for all cases

        """

        nyha_class = values[self.NYHA].index(1) + 1

        ef = values[self.LVEF][0]

        qrs = values[self.QRS][0]

        lbbb = bool(values[self.LBBB][0])

        sr = bool(values[self.SINUS_RHYTHM][0])

        if ef > 35: #EF higher than 35 disqualified

            return 'red'

        #print "NYHA:", nyha_class, "EF:", ef, "QRS:", qrs, "LBBB:", lbbb, "SR:", sr

        if nyha_class == 1:

            if ef < 30 and qrs >= 150 and lbbb:

                return 'orange'

            else:

                return 'red'

        elif nyha_class == 2:

            if lbbb:

                if sr:

                    if qrs > 150: 

                        return 'green'

                    else: 

                        return 'yellow'

                else: 

                    return 'red'

            else:

                if qrs > 150:

                    return 'orange'

                else:

                    return 'red'

        elif nyha_class == 3:

            if sr and qrs > 120:

                if qrs > 150 and lbbb:

                    return 'green'

                elif qrs <= 150 and not lbbb:

                    return 'orange'

                else:

                    return 'yellow'

            else:

                return 'red'

        else: #this is None and class 4

            return 'red'

    def prediction_from_color(self, color):

        mapping = {'red' : 0, 'orange' : 0, 'yellow': 1, 'green' : 1}

        return mapping[color]