import pickle
from pathlib import Path
from pprint import pformat
import numpy as np
import feather
from sklearn.model_selection import StratifiedKFold
from config_tosh import conf
from utils import load_fs_tosh, now
from utils import make_submission, copy_script
# from ae_cv import AE_cv
from nn_cv import NN_cv
# from deepctr_cv import deepctr_cv

from logger import Logger


def main():
    experiment_name = now()
    cv_path = Path(f"result/{experiment_name}")
    cv_path.mkdir(parents=True)

    copy_script(cv_path)
    log = Logger(experiment_name, cv_path / "exp.log")

    log.info("load data")
    with log.interval_timer("load data"):
        train_X = load_fs_tosh('all_snap', conf)
        train_y = feather.read_dataframe("features/HasDetections.ftr")
        train_y = train_y.HasDetections
        test = load_fs_tosh('all_snap', conf, test=True)

    log.info(pformat(list(train_X.columns)))
    cv = StratifiedKFold(n_splits=5, random_state=conf.seed)
    cv = cv.split(train_X, train_y)

    log.info("learning start")
    log.double_kiritori()
    with open('features/NN/conf_tosh_all_snap.pkl', 'rb') as p:
        embedd_conf = pickle.load(p)
    log.info(pformat(embedd_conf))
    score, pred, meta = NN_cv(
        train_X, train_y, cv, log, cv_path,
        X_test=test,
        split_conf=embedd_conf
    )
    log.info(score)
    log.double_kiritori()
    log.info("done")

    del train_X, train_y

    np.save(cv_path / "test_preds.npy", pred)
    np.save(cv_path / "oof_preds.npy", meta)

    make_submission(pred, f"submissions/{experiment_name}.csv.gz")


if __name__ == "__main__":
    main()
