from sklearn.cross_validation import StratifiedKFold
from sklearn.pipeline import Pipeline
from sklearn.metrics import r2_score

from mempamal.configuration import JSONify_estimator, JSONify_cv, build_dataset
from mempamal.workflow import create_wf, save_wf
from mempamal.datasets import iris

# iris dataset as usual but with linear regression (Why not! :p).
X, y = iris.get_data()

import inspect
from mempamal.examples.elasticnet_parsimony import EnetWrap
print(inspect.getsource(inspect.getmodule(EnetWrap)))

est = Pipeline([("enet", EnetWrap())])

alphas = [1e-4, 1e-3, 1e-2, 0.1, 1., 10., 100., 1e3]
ls = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
grid = []
for a in alphas:
    for l in ls:
        grid.append({"enet__l": l,
                     "enet__alpha": a})
print("The grid contains {} sets of parameters:".format(len(grid)))
grid

method_conf = JSONify_estimator(est, out="./est.json")
cv_conf = JSONify_cv(StratifiedKFold, cv_kwargs={"n_folds": 5},
                     score_func=r2_score, 
                     stratified=True,
                     inner_cv=StratifiedKFold,
                     inner_cv_kwargs={"n_folds": 5},
                     inner_score_func=r2_score,
                     out="./cv.json")
dataset = build_dataset(X, y, method_conf, cv_conf, grid=grid, outputdir=".")
wfi = create_wf(dataset['folds'], cv_conf, method_conf, ".",
               verbose=True)
wf = save_wf(wfi, "./workflow.json", mode="soma-workflow")

from soma_workflow.client import WorkflowController

import time
import json
import sklearn.externals.joblib as joblib

controller = WorkflowController()
wf_id = controller.submit_workflow(workflow=wf, name="third example")

while controller.workflow_status(wf_id) != 'workflow_done':
    time.sleep(2)
print(joblib.load('./final_res.pkl'))