import tensorflow as tf


def pywatts_input_fn(X, y=None, num_epochs=None, shuffle=True, batch_size=400):
    return tf.estimator.inputs.pandas_input_fn(x=X,
                                               y=y,
                                               num_epochs=num_epochs,
                                               shuffle=shuffle,
                                               batch_size=batch_size)


class Net:
    __regressor = None
    __feature_cols = [tf.feature_column.numeric_column(col) for col in ['dc', 'temp', 'wind']]

    def __init__(self, feature_cols=__feature_cols):
        self.__regressor = tf.estimator.DNNRegressor(feature_columns=feature_cols,
                                                     hidden_units=[50, 50],
                                                     model_dir='tf_pywatts_model')

    def train(self, training_data, training_results, steps):
        self.__regressor.train(input_fn=pywatts_input_fn(training_data, y=training_results, num_epochs=None, shuffle=True), steps=steps)

    def evaluate(self, eval_data, eval_results):
        return self.__regressor.evaluate(input_fn=pywatts_input_fn(eval_data, y=eval_results, num_epochs=1, shuffle=False), steps=1)

    def predict1h(self, predict_data):
        return self.__regressor.predict(input_fn=pywatts_input_fn(predict_data, num_epochs=1, shuffle=False))