# Note: User can run this locally with `heroku run bash` and then `python`
# This will populate the shared database for later steps

import os
import psycopg2

import numpy as np
import pandas as pd

import statsmodels.api as sm 

DATABASE_URL = os.environ['DATABASE_URL']
conn = psycopg2.connect(DATABASE_URL, sslmode='require')
cur = conn.cursor() # cursor needed to perform db ops

cur.execute("""
    CREATE TABLE Iris (
        Sepal_Length real,
        Sepal_Width real,
        Petal_Length real,
        Petal_Width real,
        Species varchar(20)
    );
""")

iris = sm.datasets.get_rdataset("iris").data

for i in range(iris.shape[0]):
    row = iris.iloc[i]
    cur.execute("INSERT INTO Iris Values(%s, %s, %s, %s, %s)",
          (row[0], row[1], row[2], row[3], row[4]))

# Let's make sure they're there
cur.execute("SELECT * FROM Iris;")
my_result = cur.fetchall();
print(pd.DataFrame(my_result))

# Now let's define an output data set

cur.execute("""
    CREATE TABLE scores (
        score_ts timestamp,
        theta1 double precision,
        theta2 double precision,
        theta3 double precision,
        theta4 double precision
    )
""")

conn.commit()

cur.close()
conn.close()

