from neo4j import GraphDatabase, basic_auth
import pandas as pd

db_location = "bolt://localhost:7687"
username = "neo4j"
password = "qwerty"

query = """
		   MATCH (p:Person)-[l:SMOKES]->(t:Tobacco), (p)-[:LIVES]-> (s:State)
		   RETURN p.sex as sex, p.age as age, s.name as state, t.name as tobacco
		"""

def get_dataset(tx):
	db_res = tx.run(query)
	training_data = pd.DataFrame([r.values() for r in result], columns=result.keys())
	return training_data

def main():
    db = GraphDatabase.driver(db_location, auth=basic_auth(username, password))

	with db.session() as session:
		df = session.read_transaction(get_dataset)
		features_columns = ['sex', 'state', 'age']
        response_column = ['tobacco']
		X_data = training_data[df.columns]
        Y_data = training_data[response_column]

		logistic_model = LogisticRegression()
        logistic_model.fit(X_data, Y_data.values)


main()