"""
Model evaluation script

python -m adpkd_segmentation.evaluate --config path_to_config_yaml --makelinks

If using a specific GPU (e.g. device 2):
CUDA_VISIBLE_DEVICES=2 python -m evaluate --config path_to_config_yaml

The makelinks flag is needed only once to create symbolic links to the data.
"""

# %%
import argparse
import json
import os
from collections import defaultdict

import torch
import yaml
from matplotlib import pyplot as plt

from adpkd_segmentation.config.config_utils import get_object_instance
from adpkd_segmentation.data.link_data import makelinks
from adpkd_segmentation.data.data_utils import masks_to_colorimg
from adpkd_segmentation.data.data_utils import tensor_dict_to_device
from adpkd_segmentation.utils.train_utils import load_model_data


# %%
def validate(
    dataloader,
    model,
    loss_metric,
    device,
    plotting_func=None,
    plotting_dict=None,
    writer=None,
    global_step=None,
    val_metric_to_check=None,
    output_losses_list=False,
):
    all_losses_and_metrics = defaultdict(list)
    num_examples = 0
    output_example_idx = (
        hasattr(dataloader.dataset, "output_idx")
        and dataloader.dataset.output_idx
    )

    for batch_idx, output in enumerate(dataloader):
        if output_example_idx:
            x_batch, y_batch, index = output
            extra_dict = dataloader.dataset.get_extra_dict(index)
            extra_dict = tensor_dict_to_device(extra_dict, device)
        else:
            x_batch, y_batch = output
            extra_dict = None
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)
        batch_size = y_batch.size(0)
        num_examples += batch_size
        with torch.no_grad():
            y_batch_hat = model(x_batch)
            losses_and_metrics = loss_metric(y_batch_hat, y_batch, extra_dict)

            for key, value in losses_and_metrics.items():
                all_losses_and_metrics[key].append(value.item() * batch_size)

            if plotting_dict is not None and batch_idx in plotting_dict:
                # TODO: add support for softmax processing
                prediction = torch.sigmoid(y_batch_hat)
                image_idx = plotting_dict[batch_idx]
                global_im_index = batch_idx * batch_size + image_idx
                extra_dict = dataloader.dataset.get_extra_dict(
                    [global_im_index]
                )
                extra_dict = tensor_dict_to_device(extra_dict, device)
                plotting_func(
                    writer=writer,
                    batch=x_batch,
                    prediction=prediction,
                    target=y_batch,
                    global_step=global_step,
                    idx=image_idx,
                    title="val_batch_{}_image_{}".format(batch_idx, image_idx),
                )
                # check DSC metric for this image
                # `loss_metric` expects raw model outputs without the sigmoid
                im_pred = y_batch_hat[image_idx].unsqueeze(0)
                im_target_mask = y_batch[image_idx].unsqueeze(0)
                im_losses = loss_metric(im_pred, im_target_mask, extra_dict)
                writer.add_scalar(
                    "val_batch_{}_image_{}_{}".format(
                        batch_idx, image_idx, val_metric_to_check
                    ),
                    im_losses[val_metric_to_check],
                    global_step,
                )

    averaged = {}
    for key, value in all_losses_and_metrics.items():
        averaged[key] = sum(all_losses_and_metrics[key]) / num_examples

    if output_losses_list:
        return averaged, all_losses_and_metrics
    return averaged


# %%
def evaluate(config):
    model_config = config["_MODEL_CONFIG"]
    loader_to_eval = config["_LOADER_TO_EVAL"]
    dataloader_config = config[loader_to_eval]
    loss_metric_config = config["_LOSSES_METRICS_CONFIG"]
    results_path = config["_RESULTS_PATH"]
    saved_checkpoint = config["_MODEL_CHECKPOINT"]
    checkpoint_format = config["_NEW_CKP_FORMAT"]

    model = get_object_instance(model_config)()
    if saved_checkpoint is not None:
        load_model_data(saved_checkpoint, model, new_format=checkpoint_format)

    dataloader = get_object_instance(dataloader_config)()
    loss_metric = get_object_instance(loss_metric_config)()
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = model.to(device)
    model.eval()
    all_losses_and_metrics = validate(dataloader, model, loss_metric, device)

    os.makedirs(results_path)
    with open("{}/val_results.json".format(results_path), "w") as fp:
        print(all_losses_and_metrics)
        json.dump(all_losses_and_metrics, fp, indent=4)

    # plotting check
    output_example_idx = (
        hasattr(dataloader.dataset, "output_idx")
        and dataloader.dataset.output_idx
    )
    data_iter = iter(dataloader)
    if output_example_idx:
        inputs, labels, _ = next(data_iter)
    else:
        inputs, labels = next(data_iter)

    inputs = inputs.to(device)
    preds = model(inputs)
    inputs = inputs.cpu()
    preds = preds.cpu()

    plot_figure_from_batch(inputs, preds)


# %%
def plot_figure_from_batch(inputs, preds, target=None, idx=0):

    f, axarr = plt.subplots(1, 2)
    axarr[0].imshow(inputs[idx][1], cmap="gray")
    axarr[1].imshow(inputs[idx][1], cmap="gray")  # background for mask
    axarr[1].imshow(masks_to_colorimg(preds[idx]), alpha=0.5)

    return f


# %%
def quick_check(config_path, run_makelinks=False):
    if run_makelinks:
        makelinks()
    with open(config_path, "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    evaluate(config)


# %%
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config", help="YAML config path", type=str, required=True
    )
    parser.add_argument(
        "--makelinks", help="Make data links", action="store_true"
    )

    args = parser.parse_args()
    with open(args.config, "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    if args.makelinks:
        makelinks()

    evaluate(config)