from ..nodes.base_node import BaseNode
from ..signal import SignalFunctionType
import matplotlib.pyplot as plt


class CrossValidationVisualizationNode(BaseNode):
    def __init__(self):
        pass

    def signal_types(self):
        return [SignalFunctionType('visualize_cross_validation',
                                   ["*machine_cross_validation"],
                                   [])
                ]

    def run(self, name, signals, mode):
        sig = {}
        if mode == "visualize_cross_validation" and len(signals["*machine_cross_validation"]) > 0:
            fig = plt.figure()
            ax = fig.add_subplot(111)
            ax.set_title('Cross Validation')
            ax.set_ylabel(signals["*machine_cross_validation"][0][1])
            ax.set_xlabel('Machines')

            labels = [cv[0][0] for cv in signals["*machine_cross_validation"]]
            scores = [cv[2] for cv in signals["*machine_cross_validation"]]

            ax.boxplot(scores)
            plt.xticks(range(1, len(labels) + 1), labels, rotation=30, size='small')
            plt.show()
        return sig
