import numpy as np

from sklearn.base import TransformerMixin

import extract_data

import loader

class RX_Transformer(TransformerMixin):

    """

    """ 

    def __init__(self):

        pass

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

        self.all_med_classes = {}

        self.dimensionality = 0

        for empi in X:

            for med in self.get_med_classes(empi):

                if med not in self.all_med_classes:

                    self.all_med_classes[med] = self.dimensionality

                    self.dimensionality += 1

        return self

    def get_feature_names(self):

        reverse_dict = {}

        for med in self.all_med_classes.keys():

            reverse_dict[self.all_med_classes[med]] = med 

        features = []

        for i in range(self.dimensionality):

            features.append(reverse_dict[i])

        return np.asarray(features)

    def transform(self, X, **transform_params):

        transformed_X = map(self.get_med_string, X)

        return np.array(transformed_X)

    def get_med_vector(self, empi):

        vector = np.zeros(self.dimensionality)

        for med in self.get_med_classes(empi):

            vector[self.all_med_classes[med]] = 1.0

        return vector 

    def get_med_classes(self, empi):

        patient = loader.get_patient_by_EMPI(empi)

        operation_date = extract_data.get_operation_date(patient)

        medications = []

        for med in patient['Med']:

            try:

                date = parse_m_d_y(med['Medication_Date'])

                if date <= procedure_date:

                    medications.extend(med['RXNORM_CLASSES'])

            except:

                pass

        return medications