# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['index_flow', 'query_flow']

# Cell
import time
import torch

from pathlib import Path
from .loader import get_image_files, archive_loader, db_loader, treemap_loader
from .crafter import crafter, preproc
from .encoder import image_encoder, text_encoder, image_query_encoder
from .indexer import join_all, build_treemap, save_archives
from .ranker import ranker, nns_to_files

# Cell
def index_flow(path):
    root = Path(path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    filepaths = get_image_files(root)
    archive_db = {}

    archive_db, new_files = archive_loader(filepaths, root, device)
    print(f"Loaded {len(archive_db)} encodings")
    print(f"Encoding {len(new_files)} new images")

#     start_time = time.perf_counter()

    crafted_files = crafter(new_files, device)
    new_embeddings = image_encoder(crafted_files, device)

    db = join_all(archive_db, new_files, new_embeddings)
    print("Building treemap")
    t = build_treemap(db)

    print(f"Saving {len(db)} encodings")
    save_paths = save_archives(root, t, db)
#     print(f"Done in {time.perf_counter() - start_time} seconds")

    return(save_paths)

# Cell
def query_flow(path, query=None, image_query=None):
    start_time = time.time()
    print('starting timer')
    root = Path(path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print("Checking files")
    dbpath = root/'memery.pt'
    db = db_loader(dbpath, device)
    treepath = root/'memery.ann'
    treemap = treemap_loader(treepath)

    filepaths = get_image_files(root)
    if treemap == None or len(db) != len(filepaths):
        print('Indexing')
        dbpath, treepath = index_flow(root)
        treemap = treemap_loader(Path(treepath))
        db = db_loader(dbpath, device)

    print('Converting query')
    if image_query:
        img = preproc(image_query)
    if query and image_query:
        text_vec = text_encoder(query, device)
        image_vec = image_query_encoder(img, device)
        query_vec = text_vec + image_vec
    elif query:
        query_vec = text_encoder(query, device)
    elif image_query:
        query_vec = image_query_encoder(img, device)
    else:
        print('No query!')

    print(f"Searching {len(db)} images")
    indexes = ranker(query_vec, treemap)
    ranked_files = nns_to_files(db, indexes)

    print(f"Done in {time.time() - start_time} seconds")

    return(ranked_files)

