from __future__ import absolute_import, division, print_function

# TensorFlow and tf.keras
import tensorflow as tf
from tensorflow import keras

# Helper libraries
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

from utils import *

tf.enable_eager_execution()

df = pd.read_pickle('data_final.p')
# print(df.head())

BATCH_SIZE = 15
SHUFFLE_BUFFER_SIZE = 100
LEARNING_RATE = 0.001

# dataset = dataframe_to_dataset_biomes(df)
dataset_size, features, dataset = dataframe_to_dataset_temp_precip(df)
print(dataset_size)
dataset = dataset.shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE).repeat()
TRAIN_SIZE = dataset_size * 0.85
TEST_SIZE = dataset_size - TRAIN_SIZE
(training, test) = (dataset.take(TRAIN_SIZE), dataset.skip(TRAIN_SIZE))

print(training.make_one_shot_iterator().get_next())

model = keras.Sequential([
    keras.layers.Dense(32, activation=tf.nn.relu, input_shape=[features]),
    keras.layers.Dense(32, activation=tf.nn.relu),
    keras.layers.Dense(2)
])

optimizer = tf.train.AdamOptimizer(LEARNING_RATE)

model.compile(loss='mse',
              optimizer=optimizer,
              metrics=['mae'])

model.summary()

EPOCHS = 1000

history = model.fit(
    training,
    epochs=EPOCHS,
    verbose=1,
    steps_per_epoch=int(dataset_size / BATCH_SIZE)
)

# i = 0
# for feature, target in dataset:
    # print('{} => {}'.format(feature, target))

print(tf.__version__)