import logging
import os
import sys
from time import time

import numpy as np

from dipy.io.gradients import read_bvals_bvecs
from dipy.io.image import load_nifti, save_nifti
from dipy.io.stateful_tractogram import Space, StatefulTractogram
from dipy.io.streamline import load_tractogram, save_tractogram
from dipy.segment.bundles import RecoBundles
from dipy.segment.mask import median_otsu
from dipy.segment.tissue import TissueClassifierHMRF, dam_classifier
from dipy.tracking import Streamlines
from dipy.workflows.utils import handle_vol_idx
from dipy.workflows.workflow import Workflow


class MedianOtsuFlow(Workflow):
    @classmethod
    def get_short_name(cls):
        return "medotsu"

    def run(
        self,
        input_files,
        save_masked=False,
        median_radius=2,
        numpass=5,
        autocrop=False,
        vol_idx=None,
        dilate=None,
        finalize_mask=False,
        out_dir="",
        out_mask="brain_mask.nii.gz",
        out_masked="dwi_masked.nii.gz",
    ):
        """Workflow wrapping the median_otsu segmentation method.

        Applies median_otsu segmentation on each file found by 'globing'
        ``input_files`` and saves the results in a directory specified by
        ``out_dir``.

        Parameters
        ----------
        input_files : string
            Path to the input volumes. This path may contain wildcards to
            process multiple inputs at once.
        save_masked : bool, optional
            Save mask.
        median_radius : int, optional
            Radius (in voxels) of the applied median filter.
        numpass : int, optional
            Number of pass of the median filter.
        autocrop : bool, optional
            If True, the masked input_volumes will also be cropped using the
            bounding box defined by the masked data. For example, if diffusion
            images are of 1x1x1 (mm^3) or higher resolution auto-cropping could
            reduce their size in memory and speed up some of the analysis.
        vol_idx : str, optional
            1D array representing indices of ``axis=-1`` of a 4D
            `input_volume`. From the command line use something like
            '1,2,3-5,7'. This input is required for 4D volumes.
        dilate : int, optional
            number of iterations for binary dilation.
        finalize_mask : bool, optional
            Whether to remove potential holes or islands.
            Useful for solving minor errors.
        out_dir : string, optional
            Output directory.
        out_mask : string, optional
            Name of the mask volume to be saved.
        out_masked : string, optional
            Name of the masked volume to be saved.
        """
        io_it = self.get_io_iterator()
        vol_idx = handle_vol_idx(vol_idx)

        for fpath, mask_out_path, masked_out_path in io_it:
            logging.info(f"Applying median_otsu segmentation on {fpath}")

            data, affine, img = load_nifti(fpath, return_img=True)
            masked_volume, mask_volume = median_otsu(
                data,
                vol_idx=vol_idx,
                median_radius=median_radius,
                numpass=numpass,
                autocrop=autocrop,
                dilate=dilate,
                finalize_mask=finalize_mask,
            )

            save_nifti(mask_out_path, mask_volume.astype(np.float64), affine)

            logging.info(f"Mask saved as {mask_out_path}")

            if save_masked:
                save_nifti(masked_out_path, masked_volume, affine, hdr=img.header)

                logging.info(f"Masked volume saved as {masked_out_path}")

        return io_it


class RecoBundlesFlow(Workflow):
    @classmethod
    def get_short_name(cls):
        return "recobundles"

    def run(
        self,
        streamline_files,
        model_bundle_files,
        greater_than=50,
        less_than=1000000,
        no_slr=False,
        clust_thr=15.0,
        reduction_thr=15.0,
        reduction_distance="mdf",
        model_clust_thr=2.5,
        pruning_thr=8.0,
        pruning_distance="mdf",
        slr_metric="symmetric",
        slr_transform="similarity",
        slr_matrix="small",
        refine=False,
        r_reduction_thr=12.0,
        r_pruning_thr=6.0,
        no_r_slr=False,
        out_dir="",
        out_recognized_transf="recognized.trk",
        out_recognized_labels="labels.npy",
    ):
        """Recognize bundles

        See :footcite:p:`Garyfallidis2018` and :footcite:p:`Chandio2020a` for
        further details about the method.

        Parameters
        ----------
        streamline_files : string
            The path of streamline files where you want to recognize bundles.
        model_bundle_files : string
            The path of model bundle files.
        greater_than : int, optional
            Keep streamlines that have length greater than
            this value in mm.
        less_than : int, optional
            Keep streamlines have length less than this value
            in mm.
        no_slr : bool, optional
            Don't enable local Streamline-based Linear
            Registration.
        clust_thr : float, optional
            MDF distance threshold for all streamlines.
        reduction_thr : float, optional
            Reduce search space by (mm).
        reduction_distance : string, optional
            Reduction distance type can be mdf or mam.
        model_clust_thr : float, optional
            MDF distance threshold for the model bundles.
        pruning_thr : float, optional
            Pruning after matching.
        pruning_distance : string, optional
            Pruning distance type can be mdf or mam.
        slr_metric : string, optional
            Options are None, symmetric, asymmetric or diagonal.
        slr_transform : string, optional
            Transformation allowed. translation, rigid, similarity or scaling.
        slr_matrix : string, optional
            Options are 'nano', 'tiny', 'small', 'medium', 'large', 'huge'.
        refine : bool, optional
            Enable refine recognized bundle.
        r_reduction_thr : float, optional
            Refine reduce search space by (mm).
        r_pruning_thr : float, optional
            Refine pruning after matching.
        no_r_slr : bool, optional
            Don't enable Refine local Streamline-based Linear
            Registration.
        out_dir : string, optional
            Output directory.
        out_recognized_transf : string, optional
            Recognized bundle in the space of the model bundle.
        out_recognized_labels : string, optional
            Indices of recognized bundle in the original tractogram.

        References
        ----------
        .. footbibliography::

        """
        slr = not no_slr
        r_slr = not no_r_slr

        bounds = [
            (-30, 30),
            (-30, 30),
            (-30, 30),
            (-45, 45),
            (-45, 45),
            (-45, 45),
            (0.8, 1.2),
            (0.8, 1.2),
            (0.8, 1.2),
        ]

        slr_matrix = slr_matrix.lower()
        if slr_matrix == "nano":
            slr_select = (100, 100)
        if slr_matrix == "tiny":
            slr_select = (250, 250)
        if slr_matrix == "small":
            slr_select = (400, 400)
        if slr_matrix == "medium":
            slr_select = (600, 600)
        if slr_matrix == "large":
            slr_select = (800, 800)
        if slr_matrix == "huge":
            slr_select = (1200, 1200)

        slr_transform = slr_transform.lower()
        if slr_transform == "translation":
            bounds = bounds[:3]
        if slr_transform == "rigid":
            bounds = bounds[:6]
        if slr_transform == "similarity":
            bounds = bounds[:7]
        if slr_transform == "scaling":
            bounds = bounds[:9]

        logging.info("### RecoBundles ###")

        io_it = self.get_io_iterator()

        t = time()
        logging.info(streamline_files)
        input_obj = load_tractogram(streamline_files, "same", bbox_valid_check=False)
        streamlines = input_obj.streamlines

        logging.info(f" Loading time {time() - t:0.3f} sec")

        rb = RecoBundles(streamlines, greater_than=greater_than, less_than=less_than)

        for _, mb, out_rec, out_labels in io_it:
            t = time()
            logging.info(mb)
            model_bundle = load_tractogram(
                mb, "same", bbox_valid_check=False
            ).streamlines
            logging.info(f" Loading time {time() - t:0.3f} sec")
            logging.info("model file = ")
            logging.info(mb)

            recognized_bundle, labels = rb.recognize(
                model_bundle,
                model_clust_thr=model_clust_thr,
                reduction_thr=reduction_thr,
                reduction_distance=reduction_distance,
                pruning_thr=pruning_thr,
                pruning_distance=pruning_distance,
                slr=slr,
                slr_metric=slr_metric,
                slr_x0=slr_transform,
                slr_bounds=bounds,
                slr_select=slr_select,
                slr_method="L-BFGS-B",
            )

            if refine:
                if len(recognized_bundle) > 1:
                    # affine
                    x0 = np.array([0, 0, 0, 0, 0, 0, 1.0, 1.0, 1, 0, 0, 0])
                    affine_bounds = [
                        (-30, 30),
                        (-30, 30),
                        (-30, 30),
                        (-45, 45),
                        (-45, 45),
                        (-45, 45),
                        (0.8, 1.2),
                        (0.8, 1.2),
                        (0.8, 1.2),
                        (-10, 10),
                        (-10, 10),
                        (-10, 10),
                    ]

                    recognized_bundle, labels = rb.refine(
                        model_bundle,
                        recognized_bundle,
                        model_clust_thr=model_clust_thr,
                        reduction_thr=r_reduction_thr,
                        reduction_distance=reduction_distance,
                        pruning_thr=r_pruning_thr,
                        pruning_distance=pruning_distance,
                        slr=r_slr,
                        slr_metric=slr_metric,
                        slr_x0=x0,
                        slr_bounds=affine_bounds,
                        slr_select=slr_select,
                        slr_method="L-BFGS-B",
                    )

            if len(labels) > 0:
                ba, bmd = rb.evaluate_results(
                    model_bundle, recognized_bundle, slr_select
                )

                logging.info(f"Bundle adjacency Metric {ba}")
                logging.info(f"Bundle Min Distance Metric {bmd}")

            new_tractogram = StatefulTractogram(
                recognized_bundle, streamline_files, Space.RASMM
            )
            save_tractogram(new_tractogram, out_rec, bbox_valid_check=False)
            logging.info("Saving output files ...")
            np.save(out_labels, np.array(labels))
            logging.info(out_rec)
            logging.info(out_labels)


class LabelsBundlesFlow(Workflow):
    @classmethod
    def get_short_name(cls):
        return "labelsbundles"

    def run(
        self,
        streamline_files,
        labels_files,
        out_dir="",
        out_bundle="recognized_orig.trk",
    ):
        """Extract bundles using existing indices (labels)

        See :footcite:p:`Garyfallidis2018` for further details about the method.

        Parameters
        ----------
        streamline_files : string
            The path of streamline files where you want to recognize bundles.
        labels_files : string
            The path of model bundle files.
        out_dir : string, optional
            Output directory.
        out_bundle : string, optional
            Recognized bundle in the space of the model bundle.

        References
        ----------
        .. footbibliography::

        """
        logging.info("### Labels to Bundles ###")

        io_it = self.get_io_iterator()
        for f_steamlines, f_labels, out_bundle in io_it:
            logging.info(f_steamlines)
            sft = load_tractogram(f_steamlines, "same", bbox_valid_check=False)
            streamlines = sft.streamlines

            logging.info(f_labels)
            location = np.load(f_labels)
            if len(location) < 1:
                bundle = Streamlines([])
            else:
                bundle = streamlines[location]

            logging.info("Saving output files ...")
            new_sft = StatefulTractogram(bundle, sft, Space.RASMM)
            save_tractogram(new_sft, out_bundle, bbox_valid_check=False)
            logging.info(out_bundle)


class ClassifyTissueFlow(Workflow):
    @classmethod
    def get_short_name(cls):
        return "extracttissue"

    def run(
        self,
        input_files,
        bvals_file=None,
        method=None,
        wm_threshold=0.5,
        b0_threshold=50,
        low_signal_threshold=50,
        nclass=None,
        beta=0.1,
        tolerance=1e-05,
        max_iter=100,
        out_dir="",
        out_tissue="tissue_classified.nii.gz",
        out_pve="tissue_classified_pve.nii.gz",
    ):
        """Extract tissue from a volume.

        Parameters
        ----------
        input_files : string
            Path to the input volumes. This path may contain wildcards to
            process multiple inputs at once.
        bvals_file : string, optional
            Path to the b-values file. Required for 'dam' method.
        method : string, optional
            Method to use for tissue extraction. Options are:
                - 'hmrf': Markov Random Fields modeling approach.
                - 'dam': Directional Average Maps, proposed by :footcite:p:`Cheng2020`.

            'hmrf' method is recommended for T1w images, while 'dam' method is
            recommended for DWI Multishell images (single shell are not recommended).
        wm_threshold : float, optional
            The threshold below which a voxel is considered white matter. For data like
            HCP, threshold of 0.5 proves to be a good choice. For data like cfin, higher
            threshold values like 0.7 or 0.8 are more suitable. Used for 'dam' method.
        b0_threshold : float, optional
            The intensity threshold for a b=0 image. used only for 'dam' method.
        low_signal_threshold : float, optional
            The threshold below which a voxel is considered to have low signal.
            Used only for 'dam' method.
        nclass : int, optional
            Number of desired classes. Used only for 'hmrf' method.
        beta : float, optional
            Smoothing parameter, the higher this number the smoother the
            output will be. Used only for 'hmrf' method.
        tolerance : float, optional
            Value that defines the percentage of change tolerated to
            prevent the ICM loop to stop. Default is 1e-05.
            If you want tolerance check to be disabled put 'tolerance = 0'.
            Used only for 'hmrf' method.
        max_iter : int, optional
            Fixed number of desired iterations. Default is 100.
            This parameter defines the maximum number of iterations the
            algorithm will perform. The loop may terminate early if the
            change in energy sum between iterations falls below the
            threshold defined by `tolerance`. However, if `tolerance` is
            explicitly set to 0, this early stopping mechanism is disabled,
            and the algorithm will run for the specified number of
            iterations unless another stopping criterion is met.
            Used only for 'hmrf' method.
        out_dir : string, optional
            Output directory.
        out_tissue : string, optional
            Name of the tissue volume to be saved.
        out_pve : string, optional
            Name of the pve volume to be saved.

        REFERENCES
        ----------
        .. footbibliography::

        """
        io_it = self.get_io_iterator()

        if not method or method.lower() not in ["hmrf", "dam"]:
            logging.error(
                f"Unknown method '{method}' for tissue extraction. "
                "Choose '--method hmrf' (for T1w) or '--method dam' (for DWI)"
            )
            sys.exit(1)

        prefix = "t1" if method.lower() == "hmrf" else "dwi"
        for i, name in enumerate(self.flat_outputs):
            if name.endswith("tissue_classified.nii.gz"):
                self.flat_outputs[i] = name.replace(
                    "tissue_classified.nii.gz", f"{prefix}_tissue_classified.nii.gz"
                )
            if name.endswith("tissue_classified_pve.nii.gz"):
                self.flat_outputs[i] = name.replace(
                    "tissue_classified_pve.nii.gz",
                    f"{prefix}_tissue_classified_pve.nii.gz",
                )

        self.update_flat_outputs(self.flat_outputs, io_it)

        for fpath, tissue_out_path, opve in io_it:
            logging.info(f"Extracting tissue from {fpath}")

            data, affine = load_nifti(fpath)

            if method.lower() == "hmrf":
                if nclass is None:
                    logging.error(
                        "Number of classes is required for 'hmrf' method. "
                        "For example, Use '--nclass 4' to specify the number of "
                        "classes."
                    )
                    sys.exit(1)
                classifier = TissueClassifierHMRF()
                _, segmentation_final, PVE = classifier.classify(
                    data, nclass, beta, tolerance=tolerance, max_iter=max_iter
                )

                save_nifti(tissue_out_path, segmentation_final, affine)
                save_nifti(opve, PVE, affine)

            elif method.lower() == "dam":
                if bvals_file is None or not os.path.isfile(bvals_file):
                    logging.error("'--bvals filename' is required for 'dam' method")
                    sys.exit(1)

                bvals, _ = read_bvals_bvecs(bvals_file, None)
                wm_mask, gm_mask = dam_classifier(
                    data,
                    bvals,
                    wm_threshold=wm_threshold,
                    b0_threshold=b0_threshold,
                    low_signal_threshold=low_signal_threshold,
                )
                result = np.zeros(wm_mask.shape)
                result[wm_mask] = 1
                result[gm_mask] = 2
                save_nifti(tissue_out_path, result, affine)
                save_nifti(
                    opve, np.stack([wm_mask, gm_mask], axis=-1).astype(np.int32), affine
                )

            logging.info(f"Tissue saved as {tissue_out_path} and PVE as {opve}")

        return io_it
