import argparse
import gc
import os
from pathlib import Path
from queue import Queue
from threading import Thread
from typing import Union, Tuple

import nnunetv2
import numpy as np
import torch
from acvl_utils.cropping_and_padding.bounding_boxes import bounding_box_to_slice
from acvl_utils.cropping_and_padding.padding import pad_nd_image
from batchgenerators.utilities.file_and_folder_operations import load_json, join
from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO
from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
from nnunetv2.inference.sliding_window_prediction import compute_gaussian
from nnunetv2.utilities.find_class_by_name import recursive_find_python_class
from nnunetv2.utilities.helpers import empty_cache, dummy_context
from nnunetv2.utilities.label_handling.label_handling import determine_num_input_channels
from nnunetv2.utilities.plans_handling.plans_handler import PlansManager
from torch._dynamo import OptimizedModule
from torch.backends import cudnn
from tqdm import tqdm


class CustomPredictor(nnUNetPredictor):
    def predict_single_npy_array(self, input_image: np.ndarray, image_properties: dict,
                                 segmentation_previous_stage: np.ndarray = None):
        torch.set_num_threads(7)
        with torch.no_grad():
            self.network = self.network.to(self.device)
            self.network.eval()

            if self.verbose:
                print('preprocessing')
            preprocessor = self.configuration_manager.preprocessor_class(verbose=self.verbose)
            data, _, image_properties = preprocessor.run_case_npy(input_image, None, image_properties,
                                                self.plans_manager,
                                                self.configuration_manager,
                                                self.dataset_json)

            data = torch.from_numpy(data)
            del input_image
            if self.verbose:
                print('predicting')

            predicted_logits = self.predict_preprocessed_image(data)

            if self.verbose: print('Prediction done')

            segmentation = self.convert_predicted_logits_to_segmentation_with_correct_shape(predicted_logits,
                                                                                            image_properties)
        return segmentation

    def initialize_from_trained_model_folder(self, model_training_output_dir: str,
                                             use_folds: Union[Tuple[Union[int, str]], None],
                                             checkpoint_name: str = 'checkpoint_final.pth'):
        """
        This is used when making predictions with a trained model
        """
        if use_folds is None:
            use_folds = nnUNetPredictor.auto_detect_available_folds(model_training_output_dir, checkpoint_name)

        dataset_json = load_json(join(model_training_output_dir, 'dataset.json'))
        plans = load_json(join(model_training_output_dir, 'plans.json'))
        plans_manager = PlansManager(plans)

        if isinstance(use_folds, str):
            use_folds = [use_folds]

        parameters = []
        for i, f in enumerate(use_folds):
            f = int(f) if f != 'all' else f
            checkpoint = torch.load(join(model_training_output_dir, f'fold_{f}', checkpoint_name),
                                    map_location=torch.device('cpu'), weights_only=False)
            if i == 0:
                trainer_name = checkpoint['trainer_name']
                configuration_name = checkpoint['init_args']['configuration']
                inference_allowed_mirroring_axes = checkpoint['inference_allowed_mirroring_axes'] if \
                    'inference_allowed_mirroring_axes' in checkpoint.keys() else None

            parameters.append(join(model_training_output_dir, f'fold_{f}', checkpoint_name))

        configuration_manager = plans_manager.get_configuration(configuration_name)
        # restore network
        num_input_channels = determine_num_input_channels(plans_manager, configuration_manager, dataset_json)
        trainer_class = recursive_find_python_class(join(nnunetv2.__path__[0], "training", "nnUNetTrainer"),
                                                    trainer_name, 'nnunetv2.training.nnUNetTrainer')
        if trainer_class is None:
            raise RuntimeError(f'Unable to locate trainer class {trainer_name} in nnunetv2.training.nnUNetTrainer. '
                               f'Please place it there (in any .py file)!')
        network = trainer_class.build_network_architecture(
            configuration_manager.network_arch_class_name,
            configuration_manager.network_arch_init_kwargs,
            configuration_manager.network_arch_init_kwargs_req_import,
            num_input_channels,
            plans_manager.get_label_manager(dataset_json).num_segmentation_heads,
            enable_deep_supervision=False
        )

        self.plans_manager = plans_manager
        self.configuration_manager = configuration_manager
        self.list_of_parameters = parameters
        self.network = network
        self.dataset_json = dataset_json
        self.trainer_name = trainer_name
        self.allowed_mirroring_axes = inference_allowed_mirroring_axes
        self.label_manager = plans_manager.get_label_manager(dataset_json)
        if ('nnUNet_compile' in os.environ.keys()) and (os.environ['nnUNet_compile'].lower() in ('true', '1', 't')) \
                and not isinstance(self.network, OptimizedModule):
            print('Using torch.compile')
            self.network = torch.compile(self.network)

    @torch.inference_mode(mode=True)
    def predict_preprocessed_image(self, image):
        empty_cache(self.device)
        data_device = torch.device('cpu')
        predicted_logits_device = torch.device('cpu')
        gaussian_device = torch.device('cpu')
        compute_device = torch.device('cuda:0')

        data, slicer_revert_padding = pad_nd_image(image, self.configuration_manager.patch_size,
                                                   'constant', {'value': 0}, True,
                                                   None)
        del image

        slicers = self._internal_get_sliding_window_slicers(data.shape[1:])

        empty_cache(self.device)

        data = data.to(data_device)
        predicted_logits = torch.zeros((self.label_manager.num_segmentation_heads, *data.shape[1:]),
                                       dtype=torch.half,
                                       device=predicted_logits_device)
        gaussian = compute_gaussian(tuple(self.configuration_manager.patch_size), sigma_scale=1. / 8,
                                    value_scaling_factor=10,
                                    device=gaussian_device, dtype=torch.float16)

        if not self.allow_tqdm and self.verbose:
            print(f'running prediction: {len(slicers)} steps')

        for p in self.list_of_parameters:
            # network weights have to be updated outside autocast!
            # we are loading parameters on demand instead of loading them upfront. This reduces memory footprint a lot.
            # each set of parameters is only used once on the test set (one image) so run time wise this is almost the
            # same
            self.network.load_state_dict(torch.load(p, map_location=compute_device)['network_weights'], weights_only=False)
            with torch.autocast(self.device.type, enabled=True):
                for sl in tqdm(slicers, disable=not self.allow_tqdm):
                    pred = self._internal_maybe_mirror_and_predict(data[sl][None].to(compute_device))[0].to(
                        predicted_logits_device)
                    pred /= (pred.max() / 100)
                    predicted_logits[sl] += (pred * gaussian)
                del pred
        empty_cache(self.device)
        return predicted_logits

    def convert_predicted_logits_to_segmentation_with_correct_shape(self, predicted_logits, props):
        old = torch.get_num_threads()
        torch.set_num_threads(7)

        # resample to original shape
        spacing_transposed = [props['spacing'][i] for i in self.plans_manager.transpose_forward]
        current_spacing = self.configuration_manager.spacing if \
            len(self.configuration_manager.spacing) == \
            len(props['shape_after_cropping_and_before_resampling']) else \
            [spacing_transposed[0], *self.configuration_manager.spacing]
        predicted_logits = self.configuration_manager.resampling_fn_probabilities(predicted_logits,
                                                                                  props[
                                                                                      'shape_after_cropping_and_before_resampling'],
                                                                                  current_spacing,
                                                                                  [props['spacing'][i] for i in
                                                                                   self.plans_manager.transpose_forward])

        segmentation = None
        pp = None
        try:
            with torch.no_grad():
                pp = predicted_logits.to('cuda:0')
                segmentation = pp.argmax(0).cpu()
                del pp
        except RuntimeError:
            del segmentation, pp
            torch.cuda.empty_cache()
            segmentation = predicted_logits.argmax(0)
        del predicted_logits

        # segmentation may be torch.Tensor but we continue with numpy
        if isinstance(segmentation, torch.Tensor):
            segmentation = segmentation.cpu().numpy()

        # put segmentation in bbox (revert cropping)
        segmentation_reverted_cropping = np.zeros(props['shape_before_cropping'],
                                                  dtype=np.uint8 if len(
                                                      self.label_manager.foreground_labels) < 255 else np.uint16)
        slicer = bounding_box_to_slice(props['bbox_used_for_cropping'])
        segmentation_reverted_cropping[slicer] = segmentation
        del segmentation

        # revert transpose
        segmentation_reverted_cropping = segmentation_reverted_cropping.transpose(self.plans_manager.transpose_backward)
        torch.set_num_threads(old)
        return segmentation_reverted_cropping


def predict_semseg(im, prop, semseg_trained_model, semseg_folds):
    # initialize predictors
    pred_semseg = CustomPredictor(
        tile_step_size=0.5,
        use_mirroring=True,
        use_gaussian=True,
        perform_everything_on_device=False,
        allow_tqdm=True
    )
    pred_semseg.initialize_from_trained_model_folder(
        semseg_trained_model,
        use_folds=semseg_folds,
        checkpoint_name='checkpoint_final.pth'
    )

    semseg_pred = pred_semseg.predict_single_npy_array(
        im, prop, None
    )
    torch.cuda.empty_cache()
    gc.collect()
    return semseg_pred


def map_labels_to_toothfairy(predicted_seg: np.ndarray) -> np.ndarray:
    # Create an array that maps the labels directly
    max_label = 42
    mapping = np.arange(max_label + 1)

    # Define the specific remapping
    remapping = {19: 21, 20: 22, 21: 23, 22: 24, 23: 25, 24: 26, 25: 27, 26: 28,
                 27: 31, 28: 32, 29: 33, 30: 34, 31: 35, 32: 36, 33: 37, 34: 38,
                 35: 41, 36: 42, 37: 43, 38: 44, 39: 45, 40: 46, 41: 47, 42: 48}

    # Apply the remapping
    for k, v in remapping.items():
        mapping[k] = v

    return mapping[predicted_seg]


def postprocess(prediction_npy, vol_per_voxel, verbose: bool = False):
    cutoffs = {1: 0.0,
               2: 78411.5,
               3: 0.0,
               4: 0.0,
               5: 2800.0,
               6: 1216.5,
               7: 0.0,
               8: 6222.0,
               9: 1573.0,
               10: 946.0,
               11: 0.0,
               12: 6783.5,
               13: 9469.5,
               14: 0.0,
               15: 2260.0,
               16: 3566.0,
               17: 6321.0,
               18: 4221.5,
               19: 5829.0,
               20: 0.0,
               21: 0.0,
               22: 468.0,
               23: 1555.0,
               24: 1291.5,
               25: 2834.5,
               26: 584.5,
               27: 0.0,
               28: 0.0,
               29: 0.0,
               30: 0.0,
               31: 1935.5,
               32: 0.0,
               33: 0.0,
               34: 6140.0,
               35: 0.0,
               36: 0.0,
               37: 0.0,
               38: 2710.0,
               39: 0.0,
               40: 0.0,
               41: 0.0,
               42: 970.0}

    vol_per_voxel_cutoffs = 0.3 * 0.3 * 0.3
    for c in cutoffs.keys():
        co = cutoffs[c]
        if co > 0:
            mask = prediction_npy == c
            pred_vol = np.sum(mask) * vol_per_voxel
            if 0 < pred_vol < (co * vol_per_voxel_cutoffs):
                prediction_npy[mask] = 0
                if verbose:
                    print(
                        f'removed label {c} because predicted volume of {pred_vol} is less than the cutoff {co * vol_per_voxel_cutoffs}')
    return prediction_npy


if __name__ == '__main__':
    os.environ['nnUNet_compile'] = 'f'

    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input_folder', type=Path, default="/input/images/cbct/")
    parser.add_argument('-o', '--output_folder', type=Path, default="/output/images/oral-pharyngeal-segmentation/")
    parser.add_argument('-sem_mod', '--semseg_trained_model', type=str,
                        default="/opt/app/_trained_model/semseg_trained_model")
    parser.add_argument('--semseg_folds', type=str, nargs='+', default=[0, 1])
    args = parser.parse_args()

    args.output_folder.mkdir(exist_ok=True, parents=True)

    semseg_folds = [i if i == 'all' else int(i) for i in args.semseg_folds]
    semseg_trained_model = args.semseg_trained_model

    rw = SimpleITKIO()

    input_files = list(args.input_folder.glob('*.nii.gz')) + list(args.input_folder.glob('*.mha'))

    for input_fname in input_files:
        output_fname = args.output_folder / input_fname.name

        # we start with the instance seg because we can then start converting that while semseg is being predicted
        # load test image
        im, prop = rw.read_images([input_fname])

        with torch.no_grad():
            semseg_pred = predict_semseg(im, prop, semseg_trained_model, semseg_folds)
            torch.cuda.empty_cache()
            gc.collect()

        # now postprocess
        semseg_pred = postprocess(semseg_pred, np.prod(prop['spacing']), True)

        semseg_pred = map_labels_to_toothfairy(semseg_pred)

        # now save
        rw.write_seg(semseg_pred, output_fname, prop)
