import tensorflow as tf
from tensorflow import keras

from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

def main():
    housing = fetch_california_housing()

    X_train_full, X_test, y_train_full, y_test = \
        train_test_split( housing.data, housing.target )

    X_train, X_validation, y_train, y_validation = \
        train_test_split( X_train_full, y_train_full )

    scaler = StandardScaler()

    X_train = scaler.fit_transform( X_train )
    X_validation = scaler.transform( X_validation )
    X_test = scaler.transform( X_test )

    # How many epochs?
    epochs = 30

    model = keras.models.Sequential([
        keras.layers.Dense( epochs, activation = 'relu', 
            input_shape = X_train.shape[1:] ),
        keras.layers.Dense( 1 )
    ])

    model.compile( loss = 'mean_squared_error', optimizer = 'sgd' )

    history  = model.fit( X_train, y_train, epochs = 20,
        validation_data = (X_validation, y_validation) )

    print( '\n **Evaluate model** \n' )

    mse_test = model.evaluate( X_test, y_test )

    # How many predictions?
    n = 8

    X_new = X_test[:n]
    y_predictions = model.predict( X_new )

    print( '\n **Predictions** \n' )

    for p in y_predictions:
        print( f'{p[0]:8.4f}' )

    print( '\n **Values** \n' )

    for p in y_test[:n]:
        print( f'{p:8.4f}' )
# end of main()

if __name__ == '__main__':
    main()