import argparse
import json
import os

from CNNectome.utils import config_loader, cosem_db


def main(
    file_in,
    file_out,
    mask="volumes/masks/validation/",
    gt_version="v0004",
    training_version="v0003.2",
    crop="113",
    raw_data_path=None,
):
    db = cosem_db.MongoCosemDB(
        write_access=True, gt_version=gt_version, training_version=training_version
    )
    with open(file_in, "r") as f:
        missing = json.load(f)
    col = db.access("evaluation", (db.training_version, db.gt_version))
    predictions_to_run = set()
    for param_set in missing:
        setup = param_set[0]
        iteration = param_set[1]
        label = param_set[2]
        raw_ds = param_set[3]

        filter = {
            "setup": setup,
            "iteration": int(iteration),
            "raw_dataset": raw_ds,
            "label": label,
            "crop": crop,
        }
        print(filter)
        doc = col.find_one(filter)
        path = doc["path"]
        prediction = tuple([setup, iteration, raw_ds, path])
        predictions_to_run.add(prediction)
    predictions_to_run = list(predictions_to_run)
    predictions_to_run.sort(key=lambda param_set: int(param_set[1]))
    if raw_data_path is None:
        raw_data_path = os.path.join(
            config_loader.get_config()["organelles"]["data_path"],
            db.get_crop_by_number(crop)["parent"],
        )
    for pred in predictions_to_run:
        if not os.path.exists(file_out.format(pred[0])):
            with open(file_out.format(pred[0]), "w") as f:
                f.write("source activate cnnectome-validation-fix\n")
        with open(file_out.format(pred[0]), "a") as f:
            command = (
                f"./submit_inference.sh 10 2 jrc_hela-2 {pred[1]} --mask_ds {mask} --output_path {pred[3]}"
                f"--raw_data_path {raw_data_path} --raw_ds {pred[2]}\n"
            )
            f.write(command)


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Generate script to run missing predictions.")
    parser.add_argument(
        "in_file",
        type=str,
        help="Path to json file containing the missing predictions as tuples of (setup, iteration, label,"
        "raw_dataset). Can be generated by evaluate_new_validation_crop.py",
    )
    parser.add_argument(
        "out_file",
        type=str,
        help="Target file location to save shell script that can be run to generate all missing predictions",
    )
    parser.add_argument(
        "--mask",
        type=str,
        default=None,
        help="mask dataset to use for prediction, default is validation mask for gt_version",
    )
    parser.add_argument(
        "--gt_version",
        type=str,
        default="v0004",
        help="gt_version for evaluation database",
    )
    parser.add_argument(
        "--training_version",
        type=str,
        default="v0003.2",
        help="training_version for evaluation database",
    )
    parser.add_argument(
        "--raw_data_path",
        type=str,
        default=None,
        help="specify raw data path for prediction, defaults to what's inferrable from crop",
    )
    parser.add_argument(
        "--crop", type=str, default="113", help="crop to filter for to find output_path"
    )
    args = parser.parse_args()
    if args.mask is None:
        args.mask = f"volumes/masks/validation/{args.gt_version.lstrip('v')}"
    main(
        args.in_file,
        args.out_file,
        mask=args.mask,
        gt_version=args.gt_version,
        training_version=args.training_version,
        crop=args.crop,
        raw_data_path=args.raw_data_path,
    )
