from copy import deepcopy
import pandas as pd
import numpy as np
import json
import os


class Pipeline:
    def __init__(self, id=0):
        self.steps = []
        self.id = id

    def add_step(self, step):
        self.steps.append(step)


class Pipeline_run:
    def __init__(self, learning_job, pipeline=None):
        self.learning_job = learning_job
        self.pipeline = pipeline
        if pipeline:
            self.steps = deepcopy(pipeline.steps)
        else:
            self.steps = []
        self.fit_outputs = {}
        self.produce_outputs = {}
        self.pred_idx = 0
        self.produce_predictions = {}
        self.first_log = True

    def fit(self, X, Y):
        self.fit_outputs = {0: {0: {'learning_job': self.learning_job, 'X': X, 'Y': Y}}}
        for step in self.steps:
            data = {}
            i = 0
            for index in step.input_indices:
                data[i] = self.fit_outputs[index[0]][index[1]]
                i += 1
            step.primitive.fit(data)
            self.fit_outputs[step.index] = step.primitive.produce(data)

            if 'predictions' in self.fit_outputs[step.index][0]:
                self.fit_outputs['predictions'] = self.fit_outputs[step.index][0]['predictions']

    def produce(self, X):
        self.produce_outputs = {0: {0: {'learning_job': self.learning_job.dataset, 'X': X, 'Y': None}}}
        last_step = None
        for step in self.steps:
            data = {}
            i = 0
            for index in step.input_indices:
                data[i] = self.produce_outputs[index[0]][index[1]]
                i += 1
            self.produce_outputs[step.index] = step.primitive.produce(data)
            last_step = step
            if not list(self.produce_outputs[step.index][0]['X'].columns) == list(self.fit_outputs[step.index][0]['X'].columns):
                # print(list(self.produce_outputs[step.index][0]['X'].columns))
                raise Exception('problem'+str(list(self.produce_outputs[step.index][0]['X'].columns)))
            if not len(pd.unique(self.produce_outputs[step.index][0]['X'].columns)) == len(self.produce_outputs[step.index][0]['X'].columns) and not step.primitive.type == 'Ensemble':
                print('debug')
            if 'predictions' in self.produce_outputs[step.index][0]:
                self.produce_outputs['predictions'] = self.produce_outputs[step.index][0]['predictions']
            if 'proba_predictions' in self.produce_outputs[step.index][0]:
                self.produce_outputs['proba_predictions'] = self.produce_outputs[step.index][0]['proba_predictions']
        return

    def get_step_output(self, step, stage='fit'):
        if stage == 'fit':
            return self.fit_outputs[step]
        else:
            return self.produce_outputs[step]

    def replace_pipeline(self, pipeline):
        pass

    def can_add_step(self, step):
        data = {}
        i = 0
        for index in step.input_indices:
            data[i] = self.produce_outputs[index]
            i += 1
        return step.primitive.can_accept(data)

    def add_step(self, step):
        self.steps.append(step)
        self.refit()
        return self.fit_outputs[step.index][0]

    def refit(self):
        for step in self.steps:
            if step.index in self.fit_outputs.keys():
                continue
            data = {}
            i = 0
            for index in step.input_indices:
                data[i] = self.fit_outputs[index[0]][index[1]]
                i += 1
            step.primitive.fit(data)
            self.fit_outputs[step.index] = step.primitive.produce(data)
            if 'predictions' in self.fit_outputs[step.index][0]:
                self.fit_outputs['predictions'] = self.fit_outputs[step.index][0]['predictions']

    def is_redundant(self, idx):
        i = idx + 1
        for i in range(len(self.steps)):
            if [idx, 0] in self.steps[i].input_indices:  # Change 0 in [idx, 0] if will be more than 1 output for a primitive
                return False
        return True

    def calculate_metadata(self):
        num_primitives = [len(self.steps)]
        families = ['data preprocess', 'feature preprocess', 'feature selection', 'feature engineering', 'Classifier', 'Regressor', 'ensemble']
        num_families = [0, 0, 0, 0, 0, 0, 0]
        num_edges = 0
        for step in self.steps:
            num_families[families.index(step.primitive.type)] += 1
            num_edges += len(step.input_indices)
        num_estimators = np.sum(np.array(num_families[-3:]))
        num_preprocess = np.sum(np.array(num_families[:-3]))
        if num_preprocess > 0:
            dimentionality = num_estimators/num_preprocess
        else:
            dimentionality = -1
        return np.array(num_primitives+num_families+[num_estimators]+[num_preprocess]+[dimentionality])


    def log_to_json(self, score=0):
        curr_pipeline_dict = {}
        path = os.path.dirname(os.path.realpath(__file__)) + '/pipelines'
        if self.first_log:
            self.first_log = False
            with open(path+'/pipelines_log.yml') as json_data:
                pipelines_dict = json.load(json_data)
            with open(path+'/pipelines_run_log.yml') as json_data:
                pipelines_run_dict = json.load(json_data)
        steps_dict = dict(zip(range(1, len(self.steps)+1), self.steps))

        def arrange_step(step):
            inputs = []
            for i in range(len(step.input_indices)):
                inputs.append(step.input_indices[i][0])
            return {'index': step.index, 'inputs': inputs, 'primitive': {'name': step.primitive.name, 'hyperparams_run': step.primitive.hyperparams_run, 'description':step.primitive.description, 'type':step.primitive.type}}

        steps_dict = {k: arrange_step(stp) for k, stp in steps_dict.items()}
        id = ""
        for stp in steps_dict.values():
            id += stp['primitive']['name'] + '_'.join(str(x) for x in stp['inputs'])  # Add hyperparameters

        curr_pipeline_dict['id'] = id
        curr_pipeline_dict['steps'] = steps_dict
        if id not in pipelines_dict:
            pipelines_dict[id] = curr_pipeline_dict
            with open(path+'/pipelines_log.yml', 'w') as f:
                json.dump(pipelines_dict, f, indent=2)

        lj_dict = {}
        lj_dict['name'] = self.learning_job.name
        lj_dict['task'] = self.learning_job.task
        lj_dict['metric'] = self.learning_job.metric.name
        lj_dict['dataset'] = self.learning_job.dataset.name
        lj_dict['pipeline_id'] = id
        lj_dict['score'] = score
        lj_id = ''.join(str(x) for x in(list(lj_dict.values())))
        if lj_id not in pipelines_run_dict:
            pipelines_run_dict[lj_id] = lj_dict
            with open(path+'/pipelines_run_log.yml', 'w') as f:
                json.dump(pipelines_run_dict, f, indent=2)

    def rm_last_step(self, idx):
        if not self.steps[-1].index == idx:
            return
        self.steps = self.steps[:-1]
        if idx in self.fit_outputs:
            del self.fit_outputs[idx]
        if idx in self.produce_outputs:
            del self.produce_outputs[idx]

from copy import deepcopy
import pandas as pd
import numpy as np
import json
import os


class Pipeline:
    def __init__(self, id=0):
        self.steps = []
        self.id = id

    def add_step(self, step):
        self.steps.append(step)


class Pipeline_run:
    def __init__(self, learning_job, pipeline=None):
        self.learning_job = learning_job
        self.pipeline = pipeline
        if pipeline:
            self.steps = deepcopy(pipeline.steps)
        else:
            self.steps = []
        self.fit_outputs = {}
        self.produce_outputs = {}
        self.pred_idx = 0
        self.produce_predictions = {}
        self.first_log = True

    def fit(self, X, Y):
        self.fit_outputs = {0: {0: {'learning_job': self.learning_job, 'X': X, 'Y': Y}}}
        for step in self.steps:
            data = {}
            i = 0
            for index in step.input_indices:
                data[i] = self.fit_outputs[index[0]][index[1]]
                i += 1
            step.primitive.fit(data)
            self.fit_outputs[step.index] = step.primitive.produce(data)

            if 'predictions' in self.fit_outputs[step.index][0]:
                self.fit_outputs['predictions'] = self.fit_outputs[step.index][0]['predictions']

    def produce(self, X):
        self.produce_outputs = {0: {0: {'learning_job': self.learning_job.dataset, 'X': X, 'Y': None}}}
        last_step = None
        for step in self.steps:
            data = {}
            i = 0
            for index in step.input_indices:
                data[i] = self.produce_outputs[index[0]][index[1]]
                i += 1
            self.produce_outputs[step.index] = step.primitive.produce(data)
            last_step = step
            if not list(self.produce_outputs[step.index][0]['X'].columns) == list(self.fit_outputs[step.index][0]['X'].columns):
                # print(list(self.produce_outputs[step.index][0]['X'].columns))
                raise Exception('problem'+str(list(self.produce_outputs[step.index][0]['X'].columns)))
            if not len(pd.unique(self.produce_outputs[step.index][0]['X'].columns)) == len(self.produce_outputs[step.index][0]['X'].columns) and not step.primitive.type == 'Ensemble':
                print('debug')
            if 'predictions' in self.produce_outputs[step.index][0]:
                self.produce_outputs['predictions'] = self.produce_outputs[step.index][0]['predictions']
            if 'proba_predictions' in self.produce_outputs[step.index][0]:
                self.produce_outputs['proba_predictions'] = self.produce_outputs[step.index][0]['proba_predictions']
        return

    def get_step_output(self, step, stage='fit'):
        if stage == 'fit':
            return self.fit_outputs[step]
        else:
            return self.produce_outputs[step]

    def replace_pipeline(self, pipeline):
        pass

    def can_add_step(self, step):
        data = {}
        i = 0
        for index in step.input_indices:
            data[i] = self.produce_outputs[index]
            i += 1
        return step.primitive.can_accept(data)

    def add_step(self, step):
        self.steps.append(step)
        self.refit()
        return self.fit_outputs[step.index][0]

    def refit(self):
        for step in self.steps:
            if step.index in self.fit_outputs.keys():
                continue
            data = {}
            i = 0
            for index in step.input_indices:
                data[i] = self.fit_outputs[index[0]][index[1]]
                i += 1
            step.primitive.fit(data)
            self.fit_outputs[step.index] = step.primitive.produce(data)
            if 'predictions' in self.fit_outputs[step.index][0]:
                self.fit_outputs['predictions'] = self.fit_outputs[step.index][0]['predictions']

    def is_redundant(self, idx):
        i = idx + 1
        for i in range(len(self.steps)):
            if [idx, 0] in self.steps[i].input_indices:  # Change 0 in [idx, 0] if will be more than 1 output for a primitive
                return False
        return True

    def calculate_metadata(self):
        num_primitives = [len(self.steps)]
        families = ['data preprocess', 'feature preprocess', 'feature selection', 'feature engineering', 'Classifier', 'Regressor', 'ensemble']
        num_families = [0, 0, 0, 0, 0, 0, 0]
        num_edges = 0
        for step in self.steps:
            num_families[families.index(step.primitive.type)] += 1
            num_edges += len(step.input_indices)
        num_estimators = np.sum(np.array(num_families[-3:]))
        num_preprocess = np.sum(np.array(num_families[:-3]))
        if num_preprocess > 0:
            dimentionality = num_estimators/num_preprocess
        else:
            dimentionality = -1
        return np.array(num_primitives+num_families+[num_estimators]+[num_preprocess]+[dimentionality])


    def log_to_json(self, score=0):
        curr_pipeline_dict = {}
        path = os.path.dirname(os.path.realpath(__file__)) + '/pipelines'
        if self.first_log:
            self.first_log = False
            with open(path+'/pipelines_log.yml') as json_data:
                pipelines_dict = json.load(json_data)
            with open(path+'/pipelines_run_log.yml') as json_data:
                pipelines_run_dict = json.load(json_data)
        steps_dict = dict(zip(range(1, len(self.steps)+1), self.steps))

        def arrange_step(step):
            inputs = []
            for i in range(len(step.input_indices)):
                inputs.append(step.input_indices[i][0])
            return {'index': step.index, 'inputs': inputs, 'primitive': {'name': step.primitive.name, 'hyperparams_run': step.primitive.hyperparams_run, 'description':step.primitive.description, 'type':step.primitive.type}}

        steps_dict = {k: arrange_step(stp) for k, stp in steps_dict.items()}
        id = ""
        for stp in steps_dict.values():
            id += stp['primitive']['name'] + '_'.join(str(x) for x in stp['inputs'])  # Add hyperparameters

        curr_pipeline_dict['id'] = id
        curr_pipeline_dict['steps'] = steps_dict
        if id not in pipelines_dict:
            pipelines_dict[id] = curr_pipeline_dict
            with open(path+'/pipelines_log.yml', 'w') as f:
                json.dump(pipelines_dict, f, indent=2)

        lj_dict = {}
        lj_dict['name'] = self.learning_job.name
        lj_dict['task'] = self.learning_job.task
        lj_dict['metric'] = self.learning_job.metric.name
        lj_dict['dataset'] = self.learning_job.dataset.name
        lj_dict['pipeline_id'] = id
        lj_dict['score'] = score
        lj_id = ''.join(str(x) for x in(list(lj_dict.values())))
        if lj_id not in pipelines_run_dict:
            pipelines_run_dict[lj_id] = lj_dict
            with open(path+'/pipelines_run_log.yml', 'w') as f:
                json.dump(pipelines_run_dict, f, indent=2)

    def rm_last_step(self, idx):
        if not self.steps[-1].index == idx:
            return
        self.steps = self.steps[:-1]
        if idx in self.fit_outputs:
            del self.fit_outputs[idx]
        if idx in self.produce_outputs:
            del self.produce_outputs[idx]

# print(list(self.produce_outputs[step.index][0]['X'].columns)) # Change 0 in [idx, 0] if will be more than 1 output for a primitive # Add hyperparameters