from sqlalchemy import Table, Column, Float, Integer, create_engine
from fireworks.extensions.database import create_table, TablePipe, DBPipe
from fireworks.extensions import SQLFactory
from fireworks.extensions import Experiment
import os

from model_selection import loss, train_set, test_set, params, get_trainer, Parameterizer, AccuracyMetric

if __name__=="__main__":

    description = "Model selection for nonlinear regression. We are comparing the regression accuracy of different polynomial models fit to data generated by a random polynomial."
    experiment = Experiment("model_selection_db", db_path=".", description=description)
    # SQL factory
    params_table = create_table('params', columns=[
        Column('a', Integer), Column('b', Integer), Column('c', Integer), Column('d', Integer), Column('e', Integer)
        ])
    metrics_tables = {'accuracy': create_table('accuracy', columns=[Column('average-loss', Float)])}
    # engine = create_engine('sqlite:///model_selection.sqlite')
    engine = experiment.get_engine('factory.sqlite')
    factory = SQLFactory(components={
        'trainer': get_trainer(train_set, loss, optimizer='Adam', lr=.1),
        'eval_set': test_set,
        'parameterizer': Parameterizer(),
        'metrics': {'accuracy': AccuracyMetric()},
        'engine': engine,
        'params_table': params_table,
        'metrics_tables': metrics_tables,
        })

    factory.run()

    params_table = DBPipe('params', factory.engine)
    print(params_table.all())
    accuracy_table = DBPipe('accuracy', factory.engine)
    print(accuracy_table.all())
