import pandas as pd
from dame_flame.flame_db.utils import *
from dame_flame.flame_db.FLAME_db_algorithm import *

train_df = pd.DataFrame([[0,1,1,1,0,5], [0,1,1,0,0,6], [1,0,1,1,1,7], [1,1,1,1,1,7]], 
                  columns=["x1", "x2", "x3", "x4", "treated", "outcome"])
test_df = pd.DataFrame([[0,1,1,1,0,5], [0,1,1,0,0,6], [1,0,1,1,1,7], [1,1,1,1,1,7]], 
                  columns=["x1", "x2", "x3", "x4", "treated", "outcome"])                 

select_db = "postgreSQL"  # Select the database you are using: "MySQL", "postgreSQL","Microsoft SQL server"
database_name='tmp' # database name you use 
host = 'localhost' 
port = "5432"
user="postgres"
password= ""

conn = connect_db(database_name, user, password, host, port)

from dame_flame.flame_db.gen_insert_data import *
insert_data_to_db("datasetToBeMatched", # The name of your table containing the dataset to be matched
                  test_df,
                  treatment_column_name= "treated",
                  outcome_column_name= 'outcome',conn = conn)

res = FLAME_db(input_data = "datasetToBeMatched", # The name of your table containing the dataset to be matched
              holdout_data = train_df, # holdout set. We will use holdout set to train our model
              conn = conn # connector object that connects to your database. This is the output from function connect_db.
              )

res[0]:
    data frame of matched groups. Each row represent one matched groups.
    res[0]['avg_outcome_control']: 
        average of control units' outcomes in each matched group   
    res[0]['avg_outcome_treated']: 
        average of treated units' outcomes in each matched group   
    res[0]['num_control']:
        the number of control units in each matched group
    res[0]['num_treated']:
        the number of treated units in each matched group
    res[0]['is_matched']:
        the level each matched group belongs to
res[1]:
    a list of level numbers where we have matched groups
res[2]:
    a list of covariate names that we dropped

ATE_db(res) # Get ATE for the whole dataset
ATT_db(res) # Get ATT for the whole dataset