# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/apply.ipynb (unless otherwise specified).

__all__ = ['preprocess_db_intra_image', 'aggregate_fold_stats']

# Cell
from ifcimglib import imglmdb, utils, preprocessing, cif2lmdb
import numpy
import matplotlib.pyplot as plt
from tqdm import trange
import pickle
import logging
import lmdb
from pathlib import Path
from tqdm import tqdm
import os
import seaborn
from sklearn.model_selection import PredefinedSplit

# Cell

def preprocess_db_intra_image(db, preprocessed_output_path):
    logger = logging.getLogger(__name__)

    if Path(preprocessed_output_path).exists():
        Path(preprocessed_output_path).unlink()

    env = lmdb.open(preprocessed_output_path, lock=False, sync=False, map_size=cif2lmdb.map_size, subdir=False)
    logger.info("Opened lmdb database %s" % preprocessed_output_path)

    with env.begin(write=True) as txn:
        txn.put(b'__targets__', pickle.dumps(db.targets))
        txn.put(b'__len__', len(db).to_bytes(db.idx_byte_length, "big"))
        txn.put(b'__names__', " ".join(db.names).encode("utf-8"))

        for i in trange(len(db)):
            x, m, _ = db.get_image(i)
            x = x.astype(numpy.float32)
            x = preprocessing.log_transform(x, m, [1])
            x = preprocessing.min_max_normalize(x, m, "clip")
            x = preprocessing.crop_and_pad_to_square(x, 70)
            m = preprocessing.crop_and_pad_to_square(m.astype(numpy.uint8), 70).astype(bool)

            instance = cif2lmdb.get_instance(x.shape[1:], x.shape[0])
            instance = cif2lmdb.set_instance_data(instance, x.astype(numpy.float16), m)

            txn.put(i.to_bytes(db.idx_byte_length, byteorder='big'), pickle.dumps(instance))
    env.sync()
    env.close()

# Cell
def aggregate_fold_stats(db_paths, cv_pkl_file):
    preprocessed_db = imglmdb.multidbwrapper(sorted(db_paths))
    with open(cv_pkl_file, "rb") as pkl:
        test_fold, nested_test_folds = pickle.load(pkl)

    splitter = PredefinedSplit(test_fold)

    data = [{}]*splitter.get_n_splits()

    for i, (nested_test_fold, (_, test_idx)) in enumerate(zip(nested_test_folds, splitter.split())):
        per_pixel_stats = preprocessing.compute_per_pixel_stats(preprocessed_db, None, idx=test_idx)
        std_per_pixel = numpy.where(per_pixel_stats[1] == 0.0, 1, per_pixel_stats[1])
        data[i]["outer"] = (per_pixel_stats[0], std_per_pixel)

        nested_splitter = PredefinedSplit(nested_test_fold)
        data[i]["nested"] = [{}]*nested_splitter.get_n_splits()

        for j, (train_idx, val_idx) in enumerate(nested_splitter.split()):
            per_pixel_stats = preprocessing.compute_per_pixel_stats(preprocessed_db, None, idx=train_idx)
            std_per_pixel = numpy.where(per_pixel_stats[1] == 0.0, 1, per_pixel_stats[1])
            data[i]["nested"][j]["train"] = (per_pixel_stats[0], std_per_pixel)

            per_pixel_stats = preprocessing.compute_per_pixel_stats(preprocessed_db, None, idx=val_idx)
            std_per_pixel = numpy.where(per_pixel_stats[1] == 0.0, 1, per_pixel_stats[1])
            data[i]["nested"][j]["val"] = (per_pixel_stats[0], std_per_pixel)

    with open(os.path.splitext(cv_pkl_file)[0] + "_stats.pkl", "wb") as pkl:
        pickle.dump(data, pkl)

    return data