import os

import logging

from glob import glob

import numpy as np

from skimage.color import label2rgb

from skimage import img_as_ubyte

from tqdm import tqdm

import click

import cv2

import tifffile

import torch

import torch.nn as nn

from torch.utils.data.dataloader import DataLoader

from rocaseg.datasets import sources_from_path

from rocaseg.components import CheckpointHandler

from rocaseg.components.formats import numpy_to_nifti, png_to_numpy

from rocaseg.models import dict_models

from rocaseg.preproc import *

from rocaseg.repro import set_ultimate_seed

# The fix is a workaround to PyTorch multiprocessing issue:

# "RuntimeError: received 0 items of ancdata"

torch.multiprocessing.set_sharing_strategy('file_system')

cv2.ocl.setUseOpenCL(False)

cv2.setNumThreads(0)

logging.basicConfig()

logger = logging.getLogger('eval')

logger.setLevel(logging.INFO)

set_ultimate_seed()

if torch.cuda.is_available():

    maybe_gpu = 'cuda'

else:

    maybe_gpu = 'cpu'

def predict_folds(config, loader, fold_idcs):

    """Evaluate the model versus each fold

    """

    for fold_idx in fold_idcs:

        paths_weights_fold = dict()

        paths_weights_fold['segm'] = \

            os.path.join(config['path_weights'], 'segm', f'fold_{fold_idx}')

        handlers_ckpt = dict()

        handlers_ckpt['segm'] = CheckpointHandler(paths_weights_fold['segm'])

        paths_ckpt_sel = dict()

        paths_ckpt_sel['segm'] = handlers_ckpt['segm'].get_last_ckpt()

        # Initialize and configure the model

        model = (dict_models[config['model_segm']]

                 (input_channels=config['input_channels'],

                  output_channels=config['output_channels'],

                  center_depth=config['center_depth'],

                  pretrained=config['pretrained'],

                  restore_weights=config['restore_weights'],

                  path_weights=paths_ckpt_sel['segm']))

        model = nn.DataParallel(model).to(maybe_gpu)

        model.eval()

        with tqdm(total=len(loader), desc=f'Eval, fold {fold_idx}') as prog_bar:

            for i, data_batch in enumerate(loader):

                xs, ys_true = data_batch['xs'], data_batch['ys']

                xs, ys_true = xs.to(maybe_gpu), ys_true.to(maybe_gpu)

                if config['model_segm'] == 'unet_lext':

                    ys_pred = model(xs)

                elif config['model_segm'] == 'unet_lext_aux':

                    ys_pred, _ = model(xs)

                else:

                    msg = f"Unknown model {config['model_segm']}"

                    raise ValueError(msg)

                ys_pred_softmax = nn.Softmax(dim=1)(ys_pred)

                ys_pred_softmax_np = ys_pred_softmax.detach().to('cpu').numpy()

                data_batch['pred_softmax'] = ys_pred_softmax_np

                # Rearrange the batch

                data_dicts = [{k: v[n] for k, v in data_batch.items()}

                              for n in range(len(data_batch['image']))]

                for k, data_dict in enumerate(data_dicts):

                    dir_base = os.path.join(

                        config['path_predicts'],

                        data_dict['patient'], data_dict['release'], data_dict['sequence'])

                    fname_base = os.path.splitext(

                        os.path.basename(data_dict['path_rel_image']))[0]

                    # Save the predictions

                    dir_predicts = os.path.join(dir_base, 'mask_folds')

                    if not os.path.exists(dir_predicts):

                        os.makedirs(dir_predicts)

                    fname_full = os.path.join(

                        dir_predicts,

                        f'{fname_base}_fold_{fold_idx}.tiff')

                    tmp = (data_dict['pred_softmax'] * 255).astype(np.uint8, casting='unsafe')

                    tifffile.imsave(fname_full, tmp, compress=9)

                prog_bar.update(1)

def merge_predictions(config, source, loader, dict_fns,

                      save_plots=False, remove_foldw=False, convert_to_nifti=True):

    """Merge the predictions over all folds

    """

    dir_source_root = source['path_root']

    df_meta = loader.dataset.df_meta

    with tqdm(total=len(df_meta), desc='Merge') as prog_bar:

        for i, row in df_meta.iterrows():

            dir_scan_predicts = os.path.join(

                config['path_predicts'],

                row['patient'], row['release'], row['sequence'])

            dir_image_prep = os.path.join(dir_scan_predicts, 'image_prep')

            dir_mask_prep = os.path.join(dir_scan_predicts, 'mask_prep')

            dir_mask_folds = os.path.join(dir_scan_predicts, 'mask_folds')

            dir_mask_foldavg = os.path.join(dir_scan_predicts, 'mask_foldavg')

            dir_vis_foldavg = os.path.join(dir_scan_predicts, 'vis_foldavg')

            for p in (dir_image_prep, dir_mask_prep, dir_mask_folds, dir_mask_foldavg,

                      dir_vis_foldavg):

                if not os.path.exists(p):

                    os.makedirs(p)

            # Find the corresponding prediction files

            fname_base = os.path.splitext(os.path.basename(row['path_rel_image']))[0]

            fnames_pred = glob(os.path.join(dir_mask_folds, f'{fname_base}_fold_*.*'))

            # Read the reference data

            image = cv2.imread(

                os.path.join(dir_source_root, row['path_rel_image']),

                cv2.IMREAD_GRAYSCALE)

            image = dict_fns['crop'](image[None, ])[0]

            image = np.squeeze(image)

            if 'path_rel_mask' in row.index:

                ys_true = loader.dataset.read_mask(

                    os.path.join(dir_source_root, row['path_rel_mask']))

                if ys_true is not None:

                    ys_true = dict_fns['crop'](ys_true)[0]

            else:

                ys_true = None

            # Read the fold-wise predictions

            yss_pred = [tifffile.imread(f) for f in fnames_pred]

            ys_pred = np.stack(yss_pred, axis=0).astype(np.float32) / 255

            ys_pred = torch.from_numpy(ys_pred).unsqueeze(dim=0)

            # Average the fold predictions

            ys_pred = torch.mean(ys_pred, dim=1, keepdim=False)

            ys_pred_softmax = ys_pred / torch.sum(ys_pred, dim=1, keepdim=True)

            ys_pred_softmax_np = ys_pred_softmax.squeeze().numpy()

            ys_pred_arg_np = ys_pred_softmax_np.argmax(axis=0)

            # Save preprocessed input data

            fname_full = os.path.join(dir_image_prep, f'{fname_base}.png')

            cv2.imwrite(fname_full, image)  # image

            if ys_true is not None:

                ys_true = ys_true.astype(np.float32)

                ys_true = torch.from_numpy(ys_true).unsqueeze(dim=0)

                ys_true_arg_np = ys_true.numpy().squeeze().argmax(axis=0)

                fname_full = os.path.join(dir_mask_prep, f'{fname_base}.png')

                cv2.imwrite(fname_full, ys_true_arg_np)  # mask

            fname_meta = os.path.join(config['path_predicts'], 'meta_dynamic.csv')

            if not os.path.exists(fname_meta):

                df_meta.to_csv(fname_meta, index=False)  # metainfo

            # Save ensemble prediction

            fname_full = os.path.join(dir_mask_foldavg, f'{fname_base}.png')

            cv2.imwrite(fname_full, ys_pred_arg_np)

            # Save ensemble visualizations

            if save_plots:

                if ys_true is not None:

                    fname_full = os.path.join(

                        dir_vis_foldavg, f"{fname_base}_overlay_mask.png")

                    save_vis_overlay(image=image,

                                     mask=ys_true_arg_np,

                                     num_classes=config['output_channels'],

                                     fname=fname_full)

                fname_full = os.path.join(

                    dir_vis_foldavg, f"{fname_base}_overlay_pred.png")

                save_vis_overlay(image=image,

                                 mask=ys_pred_arg_np,

                                 num_classes=config['output_channels'],

                                 fname=fname_full)

                if ys_true is not None:

                    fname_full = os.path.join(

                        dir_vis_foldavg, f"{fname_base}_overlay_diff.png")

                    save_vis_mask_diff(image=image,

                                       mask_true=ys_true_arg_np,

                                       mask_pred=ys_pred_arg_np,

                                       fname=fname_full)

            # Remove the fold predictions

            if remove_foldw:

                for f in fnames_pred:

                    try:

                        os.remove(f)

                    except OSError:

                        logger.error(f'Cannot remove {f}')

            prog_bar.update(1)

    # Convert the results to 3D NIfTI images

    if convert_to_nifti:

        df_meta = df_meta.sort_values(by=["patient", "release", "sequence", "side"])

        for gb_name, gb_df in tqdm(

                df_meta.groupby(["patient", "release", "sequence", "side"]),

                desc="Convert to NIfTI"):

            patient, release, sequence, side = gb_name

            spacings = (gb_df['pixel_spacing_0'].iloc[0],

                        gb_df['pixel_spacing_1'].iloc[0],

                        gb_df['slice_thickness'].iloc[0])

            dir_scan_predicts = os.path.join(config['path_predicts'],

                                             patient, release, sequence)

            for result in ("image_prep", "mask_prep", "mask_foldavg"):

                pattern = os.path.join(dir_scan_predicts, result, '*.png')

                path_nii = os.path.join(dir_scan_predicts, f"{result}.nii")

                # Read and compose 3D image

                img = png_to_numpy(pattern_fname_in=pattern, reverse=False)

                # Save to NIfTI

                numpy_to_nifti(stack=img, fname_out=path_nii,

                               spacings=spacings, rcp_to_ras=True)

def save_vis_overlay(image, mask, num_classes, fname):

    # Add a sample of each class to have consistent class colors

    mask[0, :num_classes] = list(range(num_classes))

    overlay = label2rgb(label=mask, image=image, bg_label=0,

                        colors=['orangered', 'gold', 'lime', 'fuchsia'])

    # Convert to uint8 to save space

    overlay = img_as_ubyte(overlay)

    # Save to file

    if overlay.ndim == 3:

        overlay = overlay[:, :, ::-1]

    cv2.imwrite(fname, overlay)

def save_vis_mask_diff(image, mask_true, mask_pred, fname):

    diff = np.empty_like(mask_true)

    diff[(mask_true == mask_pred) & (mask_pred == 0)] = 0  # TN

    diff[(mask_true == mask_pred) & (mask_pred != 0)] = 0  # TP

    diff[(mask_true != mask_pred) & (mask_pred == 0)] = 2  # FP

    diff[(mask_true != mask_pred) & (mask_pred != 0)] = 3  # FN

    diff_colors = ('green', 'red', 'yellow')

    diff[0, :4] = [0, 1, 2, 3]

    overlay = label2rgb(label=diff, image=image, bg_label=0,

                        colors=diff_colors)

    # Convert to uint8 to save space

    overlay = img_as_ubyte(overlay)

    # Save to file

    if overlay.ndim == 3:

        overlay = overlay[:, :, ::-1]

    cv2.imwrite(fname, overlay)

@click.command()

@click.option('--path_data_root', default='../../data')

@click.option('--path_experiment_root', default='../../results/temporary')

@click.option('--model_segm', default='unet_lext')

@click.option('--center_depth', default=1, type=int)

@click.option('--pretrained', is_flag=True)

@click.option('--restore_weights', is_flag=True)

@click.option('--input_channels', default=1, type=int)

@click.option('--output_channels', default=1, type=int)

@click.option('--dataset', type=click.Choice(

    ['oai_imo', 'okoa', 'maknee']))

@click.option('--subset',  type=click.Choice(

    ['test', 'all']))

@click.option('--mask_mode', default='all_unitibial_unimeniscus', type=str)

@click.option('--sample_mode', default='x_y', type=str)

@click.option('--batch_size', default=64, type=int)

@click.option('--fold_num', default=5, type=int)

@click.option('--fold_idx', default=-1, type=int)

@click.option('--fold_idx_ignore', multiple=True, type=int)

@click.option('--num_workers', default=1, type=int)

@click.option('--seed_trainval_test', default=0, type=int)

@click.option('--predict_folds', is_flag=True)

@click.option('--merge_predictions', is_flag=True)

@click.option('--save_plots', is_flag=True)

def main(**config):

    config['path_data_root'] = os.path.abspath(config['path_data_root'])

    config['path_experiment_root'] = os.path.abspath(config['path_experiment_root'])

    config['path_weights'] = os.path.join(config['path_experiment_root'], 'weights')

    if not os.path.exists(config['path_weights']):

        raise ValueError('{} does not exist'.format(config['path_weights']))

    config['path_predicts'] = os.path.join(

        config['path_experiment_root'], f"predicts_{config['dataset']}_test")

    config['path_logs'] = os.path.join(

        config['path_experiment_root'], f"logs_{config['dataset']}_test")

    os.makedirs(config['path_predicts'], exist_ok=True)

    os.makedirs(config['path_logs'], exist_ok=True)

    logging_fh = logging.FileHandler(

        os.path.join(config['path_logs'], 'main.log'))

    logging_fh.setLevel(logging.DEBUG)

    logger.addHandler(logging_fh)

    # Collect the available and specified sources

    sources = sources_from_path(path_data_root=config['path_data_root'],

                                selection=config['dataset'],

                                with_folds=True,

                                seed_trainval_test=config['seed_trainval_test'])

    # Select the subset for evaluation

    if config['subset'] == 'test':

        logging.warning('Using the regular trainval-test split')

    elif config['subset'] == 'all':

        logging.warning('Using data selection: full dataset')

        for s in sources:

            sources[s]['test_df'] = sources[s]['sel_df']

            logger.info(f"Selected number of samples: {len(sources[s]['test_df'])}")

    else:

        raise ValueError(f"Unknown dataset: {config['subset']}")

    if config['dataset'] == 'oai_imo':

        from rocaseg.datasets import DatasetOAIiMoSagittal2d as DatasetSagittal2d

    elif config['dataset'] == 'okoa':

        from rocaseg.datasets import DatasetOKOASagittal2d as DatasetSagittal2d

    elif config['dataset'] == 'maknee':

        from rocaseg.datasets import DatasetMAKNEESagittal2d as DatasetSagittal2d

    else:

        raise ValueError(f"Unknown dataset: {config['dataset']}")

    # Configure dataset-dependent transforms

    fn_crop = CenterCrop(height=300, width=300)

    if config['dataset'] == 'oai_imo':

        fn_norm = Normalize(mean=0.252699, std=0.251142)

        fn_unnorm = UnNormalize(mean=0.252699, std=0.251142)

    elif config['dataset'] == 'okoa':

        fn_norm = Normalize(mean=0.232454, std=0.236259)

        fn_unnorm = UnNormalize(mean=0.232454, std=0.236259)

    else:

        msg = f"No transforms defined for dataset: {config['dataset']}"

        raise NotImplementedError(msg)

    dict_fns = {'crop': fn_crop, 'norm': fn_norm, 'unnorm': fn_unnorm}

    dataset_test = DatasetSagittal2d(

        df_meta=sources[config['dataset']]['test_df'], mask_mode=config['mask_mode'],

        name=config['dataset'], sample_mode=config['sample_mode'],

        transforms=[

            PercentileClippingAndToFloat(cut_min=10, cut_max=99),

            fn_crop,

            fn_norm,

            ToTensor()

        ])

    loader_test = DataLoader(dataset_test,

                             batch_size=config['batch_size'],

                             shuffle=False,

                             num_workers=config['num_workers'],

                             drop_last=False)

    # Build a list of folds to run on

    if config['fold_idx'] == -1:

        fold_idcs = list(range(config['fold_num']))

    else:

        fold_idcs = [config['fold_idx'], ]

    for g in config['fold_idx_ignore']:

        fold_idcs = [i for i in fold_idcs if i != g]

    # Execute

    with torch.no_grad():

        if config['predict_folds']:

            predict_folds(config=config, loader=loader_test, fold_idcs=fold_idcs)

        if config['merge_predictions']:

            merge_predictions(config=config, source=sources[config['dataset']],

                              loader=loader_test, dict_fns=dict_fns,

                              save_plots=config['save_plots'], remove_foldw=False,

                              convert_to_nifti=True)

if __name__ == '__main__':

    main()