# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""ReportCapableInterfaces for segmentation tools."""
import os

from nipype.interfaces.base import File, isdefined
from nipype.interfaces import fsl, freesurfer
from nipype.interfaces.mixins import reporting
from ... import NIWORKFLOWS_LOG
from . import base as nrb


class _FASTInputSpecRPT(nrb._SVGReportCapableInputSpec, fsl.preprocess.FASTInputSpec):
    pass


class _FASTOutputSpecRPT(
    reporting.ReportCapableOutputSpec, fsl.preprocess.FASTOutputSpec
):
    pass


class FASTRPT(nrb.SegmentationRC, fsl.FAST):
    input_spec = _FASTInputSpecRPT
    output_spec = _FASTOutputSpecRPT

    def _run_interface(self, runtime):
        if self.generate_report:
            self.inputs.segments = True

        return super(FASTRPT, self)._run_interface(runtime)

    def _post_run_hook(self, runtime):
        """generates a report showing nine slices, three per axis, of an
        arbitrary volume of `in_files`, with the resulting segmentation
        overlaid"""
        self._anat_file = self.inputs.in_files[0]
        outputs = self.aggregate_outputs(runtime=runtime)
        self._mask_file = outputs.tissue_class_map
        # We are skipping the CSF class because with combination with others
        # it only shows the skullstriping mask
        self._seg_files = outputs.tissue_class_files[1:]
        self._masked = False

        NIWORKFLOWS_LOG.info(
            "Generating report for FAST (in_files %s, "
            "segmentation %s, individual tissue classes %s).",
            self.inputs.in_files,
            outputs.tissue_class_map,
            outputs.tissue_class_files,
        )

        return super(FASTRPT, self)._post_run_hook(runtime)


class _ReconAllInputSpecRPT(
    nrb._SVGReportCapableInputSpec, freesurfer.preprocess.ReconAllInputSpec
):
    pass


class _ReconAllOutputSpecRPT(
    reporting.ReportCapableOutputSpec, freesurfer.preprocess.ReconAllOutputSpec
):
    pass


class ReconAllRPT(nrb.SurfaceSegmentationRC, freesurfer.preprocess.ReconAll):
    input_spec = _ReconAllInputSpecRPT
    output_spec = _ReconAllOutputSpecRPT

    def _post_run_hook(self, runtime):
        """generates a report showing nine slices, three per axis, of an
        arbitrary volume of `in_files`, with the resulting segmentation
        overlaid"""
        outputs = self.aggregate_outputs(runtime=runtime)
        self._anat_file = os.path.join(
            outputs.subjects_dir, outputs.subject_id, "mri", "brain.mgz"
        )
        self._contour = os.path.join(
            outputs.subjects_dir, outputs.subject_id, "mri", "ribbon.mgz"
        )
        self._masked = False

        NIWORKFLOWS_LOG.info(
            "Generating report for ReconAll (subject %s)", outputs.subject_id
        )

        return super(ReconAllRPT, self)._post_run_hook(runtime)


class _MELODICInputSpecRPT(nrb._SVGReportCapableInputSpec, fsl.model.MELODICInputSpec):
    out_report = File(
        "melodic_reportlet.svg",
        usedefault=True,
        desc="Filename for the visual" " report generated " "by Nipype.",
    )
    report_mask = File(
        desc="Mask used to draw the outline on the reportlet. "
        "If not set the mask will be derived from the data."
    )


class _MELODICOutputSpecRPT(
    reporting.ReportCapableOutputSpec, fsl.model.MELODICOutputSpec
):
    pass


class MELODICRPT(fsl.MELODIC):
    input_spec = _MELODICInputSpecRPT
    output_spec = _MELODICOutputSpecRPT
    _out_report = None

    def __init__(self, generate_report=False, **kwargs):
        super(MELODICRPT, self).__init__(**kwargs)
        self.generate_report = generate_report

    def _post_run_hook(self, runtime):
        # Run _post_run_hook of super class
        runtime = super(MELODICRPT, self)._post_run_hook(runtime)
        # leave early if there's nothing to do
        if not self.generate_report:
            return runtime

        NIWORKFLOWS_LOG.info("Generating report for MELODIC.")
        _melodic_dir = runtime.cwd
        if isdefined(self.inputs.out_dir):
            _melodic_dir = self.inputs.out_dir
        self._melodic_dir = os.path.abspath(_melodic_dir)

        self._out_report = self.inputs.out_report
        if not os.path.isabs(self._out_report):
            self._out_report = os.path.abspath(
                os.path.join(runtime.cwd, self._out_report)
            )

        mix = os.path.join(self._melodic_dir, "melodic_mix")
        if not os.path.exists(mix):
            NIWORKFLOWS_LOG.warning(
                "MELODIC outputs not found, assuming it didn't converge."
            )
            self._out_report = self._out_report.replace(".svg", ".html")
            snippet = "<h4>MELODIC did not converge, no output</h4>"
            with open(self._out_report, "w") as fobj:
                fobj.write(snippet)
            return runtime

        self._generate_report()
        return runtime

    def _list_outputs(self):
        try:
            outputs = super(MELODICRPT, self)._list_outputs()
        except NotImplementedError:
            outputs = {}
        if self._out_report is not None:
            outputs["out_report"] = self._out_report
        return outputs

    def _generate_report(self):
        from niworkflows.viz.utils import plot_melodic_components

        plot_melodic_components(
            melodic_dir=self._melodic_dir,
            in_file=self.inputs.in_files[0],
            tr=self.inputs.tr_sec,
            out_file=self._out_report,
            compress=self.inputs.compress_report,
            report_mask=self.inputs.report_mask,
        )


class _ICA_AROMAInputSpecRPT(
    nrb._SVGReportCapableInputSpec, fsl.aroma.ICA_AROMAInputSpec
):
    out_report = File(
        "ica_aroma_reportlet.svg",
        usedefault=True,
        desc="Filename for the visual" " report generated " "by Nipype.",
    )
    report_mask = File(
        desc="Mask used to draw the outline on the reportlet. "
        "If not set the mask will be derived from the data."
    )


class _ICA_AROMAOutputSpecRPT(
    reporting.ReportCapableOutputSpec, fsl.aroma.ICA_AROMAOutputSpec
):
    pass


class ICA_AROMARPT(reporting.ReportCapableInterface, fsl.ICA_AROMA):
    input_spec = _ICA_AROMAInputSpecRPT
    output_spec = _ICA_AROMAOutputSpecRPT

    def _generate_report(self):
        from niworkflows.viz.utils import plot_melodic_components

        plot_melodic_components(
            melodic_dir=self.inputs.melodic_dir,
            in_file=self.inputs.in_file,
            out_file=self.inputs.out_report,
            compress=self.inputs.compress_report,
            report_mask=self.inputs.report_mask,
            noise_components_file=self._noise_components_file,
        )

    def _post_run_hook(self, runtime):
        outputs = self.aggregate_outputs(runtime=runtime)
        self._noise_components_file = os.path.join(
            outputs.out_dir, "classified_motion_ICs.txt"
        )

        NIWORKFLOWS_LOG.info("Generating report for ICA AROMA")

        return super(ICA_AROMARPT, self)._post_run_hook(runtime)
