# coding: utf-8
# /*##########################################################################
#
# Copyright (c) 2016-2020 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/

__authors__ = ["H. Payno"]
__license__ = "MIT"
__date__ = "26/10/2021"


import copy
import functools
import logging
import os
from typing import Iterable, Optional

from nabu import version as nabu_version
from nabu.pipeline.config import generate_nabu_configfile
from nabu.pipeline.fullfield.nabu_config import (
    nabu_config as nabu_fullfield_default_config,
)
from processview.core.manager.manager import ProcessManager, DatasetState
from processview.core.superviseprocess import SuperviseProcess

from silx.io.utils import h5py_read_dataset
from silx.io.utils import open as open_hdf5
from tomwer.core.utils.deprecation import deprecated_warning

from tomwer.core.cluster.cluster import SlurmClusterConfiguration
from tomwer.core.futureobject import FutureTomwerObject
from tomwer.core.process.reconstruction.nabu.plane import NabuPlane
from tomwer.core.process.task import Task
from tomwer.core.process.icat.gallery import (
    IcatScreenshots,
    deduce_dataset_gallery_location,
    select_screenshot_from_volume,
)
from tomwer.core.scan.edfscan import EDFTomoScan
from tomwer.core.scan.nxtomoscan import NXtomoScan
from tomwer.core.scan.scanbase import TomwerScanBase
from tomwer.core.scan.scanfactory import ScanFactory
from tomwer.core.utils.dictutils import concatenate_dict
from tomwer.core.utils.scanutils import data_identifier_to_scan
from tomwer.core.volume.volumefactory import VolumeFactory
from tomwer.io.utils.h5pyutils import EntryReader
from tomwer.utils import docstring
from tomwer.io.utils import format_stderr_stdout

from . import settings, utils
from .nabucommon import ResultsLocalRun, ResultSlurmRun, _NabuBaseReconstructor
from .target import Target

_logger = logging.getLogger(__name__)


def run_volume_reconstruction(
    scan: TomwerScanBase,
    config: dict,
    dry_run: bool,
    process_id: Optional[int] = None,
) -> tuple:
    """
    Run a volume reconstruction. Scan need to have reconstruction parameters for nabu.

    Behavior: will clear link to the last volume reconstructed

    :param int process_id: optional process id

    :return: succeed, stdouts, stderrs, configs, future_scan
    :rtype: tuple
    """

    if scan.nabu_recons_params in ({}, None):
        raise ValueError(
            "no configuration provided. You should run a "
            "reconstruction from nabuslices first."
        )
    cluster_config = config.pop("cluster_config", None)
    if cluster_config == {}:
        cluster_config = None
    elif isinstance(cluster_config, SlurmClusterConfiguration):
        cluster_config = cluster_config.to_dict()

    if cluster_config is None:
        target = Target.LOCAL
    else:
        target = Target.SLURM

    # # beam shape is not directly used by nabu (uses ctf_geometry directly)
    # config.get("phase", {}).pop("beam_shape", None)

    # config_volume = copy.copy(config)
    # config_nabu_slices = copy.deepcopy(scan.nabu_recons_params)
    # if "tomwer_slices" in config_nabu_slices:
    #     del config_nabu_slices["tomwer_slices"]

    # if "phase" in config_nabu_slices and "delta_beta" in config_nabu_slices["phase"]:
    #     pag_dbs = config_nabu_slices["phase"]["delta_beta"]
    #     if isinstance(pag_dbs, str):
    #         pag_dbs = utils.retrieve_lst_of_value_from_str(pag_dbs, type_=float)
    #     if len(pag_dbs) > 1:
    #         raise ValueError(
    #             "Several value of delta / beta found for volume reconstruction"
    #         )
    # scan.clear_latest_vol_reconstructions()

    if process_id is not None:
        try:
            process_name = ProcessManager().get_process(process_id=process_id).name
        except KeyError:
            process_name = "unknow"
    else:
        process_name = ""

    volume_reconstructor = VolumeRunner(
        scan=scan,
        config_nabu=config,
        cluster_config=cluster_config,
        dry_run=dry_run,
        target=target,
        process_name=process_name,
    )
    try:
        results = volume_reconstructor.run()
    except TimeoutError as e:
        _logger.error(e)
        return None
    else:
        assert len(results) == 1, "only one volume should be reconstructed"
        res = results[0]
        # tag latest reconstructions
        if isinstance(res, ResultsLocalRun) and res.results_identifiers is not None:
            scan.set_latest_vol_reconstructions(res.results_identifiers)
        # create future if needed
        if isinstance(res, ResultSlurmRun):
            future_tomo_obj = FutureTomwerObject(
                tomo_obj=scan,
                futures=tuple(res.future_slurm_jobs),
                process_requester_id=process_id,
            )

        else:
            future_tomo_obj = None
        succeed = res.success
        stdouts = (
            [
                res.std_out,
            ]
            if hasattr(res, "std_out")
            else []
        )
        stderrs = (
            [
                res.std_err,
            ]
            if hasattr(res, "std_err")
            else []
        )
        configs = (
            [
                res.config,
            ]
            if res is not None
            else []
        )

        return succeed, stdouts, stderrs, configs, future_tomo_obj


class VolumeRunner(_NabuBaseReconstructor):
    """
    Class used to reconstruct a full volume with Nabu.
    Locally or on a cluster.
    """

    EXPECTS_SINGLE_SLICE = False

    def __init__(
        self,
        scan: TomwerScanBase,
        config_nabu,
        cluster_config: Optional[dict],
        dry_run: bool,
        target: Target,
        process_name: str,
    ) -> None:
        super().__init__(
            scan=scan,
            dry_run=dry_run,
            target=target,
            cluster_config=cluster_config,
            process_name=process_name,
        )
        self._config = config_nabu

    @property
    def configuration(self):
        return self._config

    @docstring(_NabuBaseReconstructor)
    def run(self) -> Iterable:
        dataset_params = self.scan.get_nabu_dataset_info()
        if "dataset" in self.configuration:
            dataset_params.update(self.configuration["dataset"])
        self.configuration["dataset"] = dataset_params
        self.configuration["resources"] = utils.get_nabu_resources_desc(
            scan=self.scan, workers=1, method="local"
        )

        # force overwrite results
        if "output" not in self.configuration:
            self.configuration["output"] = {}
        config_slices, cfg_folder = self._treateOutputConfig(self.configuration)

        # force overwriting results
        config_slices["output"].update({"overwrite_results": 1})

        # check and clamp `start_z` and `end_z`
        if "reconstruction" in self.configuration:
            for key in ("start_z", "end_z"):
                value = config_slices["reconstruction"].get(key)
                if value is None:
                    continue

                value = int(value)
                if self.scan.dim_2 is not None and value >= self.scan.dim_2:
                    _logger.warning(
                        f"{key} > max_size (radio height: {self.scan.dim_2}). Set it to -1 (maximum)"
                    )
                    value = -1
                config_slices["reconstruction"][key] = value

        name = (
            config_slices["output"]["file_prefix"] + settings.NABU_CONFIG_FILE_EXTENSION
        )
        if not isinstance(self.scan, EDFTomoScan):
            name = "_".join((self.scan.entry.lstrip("/"), name))
        conf_file = os.path.join(cfg_folder, name)
        _logger.info(f"{self.scan}: create {conf_file}")

        # make sure output location exists
        os.makedirs(config_slices["output"]["location"], exist_ok=True)

        # add some tomwer metadata and save the configuration
        # note: for now the section is ignored by nabu but shouldn't stay that way
        with utils.TomwerInfo(config_slices) as config_to_dump:
            generate_nabu_configfile(
                conf_file,
                nabu_fullfield_default_config,
                config=config_to_dump,
                options_level="advanced",
            )
            return tuple(
                [
                    self._process_config(
                        config_to_dump=config_to_dump,
                        config_file=conf_file,
                        info="nabu volume reconstruction",
                        file_format=config_slices["output"]["file_format"],
                        process_name=self.process_name,
                    )
                ]
            )

    @docstring(_NabuBaseReconstructor)
    def _get_futures_slurm_callback(self, config_to_dump) -> tuple:
        # add callback to set slices reconstructed urls
        class CallBack:
            # we cannot create a future directly because distributed enforce
            # the callback to have a function signature with only the future
            # as single parameter.
            def __init__(self, f_partial, scan) -> None:
                self.f_partial = f_partial
                self.scan = scan

            def process(self, fn):
                if fn.done() and not (fn.cancelled() or fn.exception()):
                    # update reconstruction urls only if processing succeed.
                    recons_identifiers = self.f_partial()
                    self.scan.add_latest_vol_reconstructions(recons_identifiers)

        file_format = config_to_dump["output"]["file_format"]
        callback = functools.partial(
            utils.get_recons_volume_identifier,
            file_prefix=config_to_dump["output"]["file_prefix"],
            location=config_to_dump["output"]["location"],
            file_format=file_format,
            scan=self.scan,
            slice_index=None,
            axis=NabuPlane.XY,  # for volume we always reconstruct along XY plane
        )

        return (CallBack(callback, self.scan),)

    def _treateOutputConfig(self, config) -> tuple:
        """

        :return: (nabu config dict, nabu extra options)
        """
        config = copy.deepcopy(config)
        config, nabu_cfg_folder = super()._treateOutputSliceConfig(config)
        os.makedirs(config["output"]["location"], exist_ok=True)

        # adapt config_s to specific volume treatment
        if "postproc" in config:
            config["postproc"] = config["postproc"]

        # make sure start_[x] and end_[x] come from config
        for key in ("start_x", "end_x", "start_y", "end_y", "start_z", "end_z"):
            if key in config:
                config["reconstruction"][key] = config[key]
                del config[key]

        return config, nabu_cfg_folder

    @docstring(_NabuBaseReconstructor)
    def _get_file_basename_reconstruction(self, pag, db, ctf, axis):
        """

        :param TomwerScanBase scan: scan reconstructed
        :param Union[None, int] slice_index: index of the slice reconstructed.
                                            if None, we want to reconstruct the
                                            entire volume
        :param bool pag: is it a paganin reconstruction
        :param int db: delta / beta parameter
        :param axis: axis over which the reconstruction goes. For volume always expected to be z. So ignored in the function
        :return: basename of the file reconstructed (without any extension)
        """
        assert type(db) in (int, type(None))
        assert not pag == ctf == True, "cannot ask for both pag and ctf active"
        if isinstance(self.scan, NXtomoScan):
            basename, _ = os.path.splitext(self.scan.master_file)
            basename = os.path.basename(basename)
            try:
                # if there is more than one entry in the file append the entry name to the file basename
                with open_hdf5(self.scan.master_file) as h5f:
                    if len(h5f.keys()) > 1:
                        basename = "_".join((basename, self.scan.entry.strip("/")))
            except Exception:
                pass
        else:
            basename = os.path.basename(self.scan.path)

        if pag:
            return "_".join((basename + "pag", "db" + str(db).zfill(4), "vol"))
        elif ctf:
            return "_".join((basename + "ctf", "db" + str(db).zfill(4), "vol"))
        else:
            return "_".join((basename, "vol"))


class NabuVolumeTask(
    Task,
    SuperviseProcess,
    input_names=("data", "nabu_params"),
    output_names=("data", "volumes", "future_tomo_obj", "screenshots"),
    optional_input_names=(
        "dry_run",
        "nabu_extra_params",  # some parameter that must update 'nabu_params' before launching the reconstruction. Such as z range...
        "serialize_output_data",
    ),
):
    def __init__(
        self,
        process_id=None,
        varinfo=None,
        inputs=None,
        node_id=None,
        node_attrs=None,
        execinfo=None,
    ):
        SuperviseProcess.__init__(self, process_id=process_id)
        Task.__init__(
            self,
            varinfo=varinfo,
            inputs=inputs,
            node_id=node_id,
            node_attrs=node_attrs,
            execinfo=execinfo,
        )
        self._dry_run = inputs.get("dry_run", False)
        self._current_processing = None

    def run(self):
        scan = data_identifier_to_scan(self.inputs.data)
        if scan is None:
            self.outputs.data = None
            return

        # update scan reconstruction parameters used
        nabu_params = concatenate_dict(
            copy.deepcopy(self.inputs.nabu_params),
            self.get_input_value("nabu_extra_params", dict()),
        )
        scan.nabu_recons_params = nabu_params
        scan.clear_latest_vol_reconstructions()

        cluster_config = nabu_params.pop("cluster_config", None)
        if cluster_config == {}:
            cluster_config = None
        elif isinstance(cluster_config, SlurmClusterConfiguration):
            cluster_config = cluster_config.to_dict()

        if "tomwer_slices" in nabu_params:
            del nabu_params["tomwer_slices"]

        if "phase" in nabu_params and "delta_beta" in nabu_params["phase"]:
            pag_dbs = nabu_params["phase"]["delta_beta"]
            if isinstance(pag_dbs, str):
                pag_dbs = utils.retrieve_lst_of_value_from_str(pag_dbs, type_=float)
            if len(pag_dbs) > 1:
                raise ValueError(
                    "Several value of delta / beta found for volume reconstruction"
                )

        if isinstance(scan, TomwerScanBase):
            scan = scan
        elif isinstance(scan, dict):
            scan = ScanFactory.create_scan_object_frm_dict(scan)
        else:
            raise ValueError(f"input type {scan} is not managed")

        if cluster_config is None:
            target = Target.LOCAL
        else:
            target = Target.SLURM

        state = None
        details = None
        try:
            self._current_processing = VolumeRunner(
                scan=scan,
                config_nabu=nabu_params,
                cluster_config=cluster_config,
                dry_run=self._dry_run,
                target=target,
                process_name=self.program_name,
            )
            res = self._current_processing.run()[0]
        except Exception as e:
            _logger.error(f"Failed to process {scan}. Error is {e}")
            state = DatasetState.FAILED
            details = None
            future_tomo_obj = None
        else:
            # tag latest reconstructions
            if isinstance(res, ResultsLocalRun) and res.results_identifiers is not None:
                scan.set_latest_vol_reconstructions(res.results_identifiers)
            # create future if needed
            if isinstance(res, ResultSlurmRun):
                future_tomo_obj = FutureTomwerObject(
                    tomo_obj=scan,
                    futures=tuple(res.future_slurm_jobs),
                    process_requester_id=self.process_id,
                )

            else:
                future_tomo_obj = None

            if self._cancelled:
                state = DatasetState.CANCELLED
                details = "cancelled by user"
                _logger.info(f"Slices computation for {scan} cancelled")
            else:
                succeed = res.success
                stdouts = (
                    [
                        res.std_out,
                    ]
                    if hasattr(res, "std_out")
                    else []
                )
                stderrs = (
                    [
                        res.std_err,
                    ]
                    if hasattr(res, "std_err")
                    else []
                )

                if not succeed:
                    mess = f"Volume computed for {scan} failed."
                    _logger.processFailed(mess)
                    state = DatasetState.FAILED
                else:
                    mess = f"Volume computed for {scan}."
                    _logger.processSucceed(mess)
                    state = DatasetState.SUCCEED

                # format stderr and stdout
                elmts = [
                    format_stderr_stdout(stderr=stderr, stdout=stdout)
                    for stderr, stdout in zip(stderrs, stdouts)
                ]
                elmts.insert(0, mess)
                details = "\n".join(elmts)
        finally:
            ProcessManager().notify_dataset_state(
                dataset=scan,
                process=self,
                state=state,
                details=details,
            )

        if self.get_input_value("serialize_output_data", True):
            self.outputs.data = scan.to_dict()
        else:
            self.outputs.data = scan
        self.outputs.volumes = scan.latest_vol_reconstructions
        self.outputs.future_tomo_obj = future_tomo_obj

        # build screenshots
        if scan.latest_vol_reconstructions is not None:
            try:
                screenshots = {}
                [
                    screenshots.update(
                        select_screenshot_from_volume(
                            VolumeFactory.create_tomo_object_from_identifier(
                                rec_identifier
                            )
                        )
                    )
                    for rec_identifier in scan.latest_vol_reconstructions
                ]
            except Exception as e:
                _logger.error(f"screenshot creation failed. Error is {e}")
            else:
                self.outputs.screenshots = IcatScreenshots(
                    data_dir=deduce_dataset_gallery_location(scan),
                    screenshots=screenshots,
                    scan=scan,
                )
        else:
            self.outputs.screenshots = tuple()

    def set_configuration(self, configuration: dict) -> None:
        Task.set_configuration(self, configuration=configuration)
        if "dry_run" in configuration:
            self.set_dry_run(bool(configuration["dry_run"]))

    @staticmethod
    def program_name():
        return "nabu-volume"

    @staticmethod
    def program_version():
        return nabu_version

    def set_dry_run(self, dry_run):
        self._dry_run = dry_run

    @property
    def dry_run(self):
        return self._dry_run

    @staticmethod
    def retrieve_last_relative_cor(scan):
        with EntryReader(scan.process_file_url) as h5f:
            latest_nabu_node = Task.get_most_recent_process(h5f, NabuVolumeTask)
            path = "configuration/reconstruction/rotation_axis_position"
            if latest_nabu_node is not None and path in latest_nabu_node:
                return h5py_read_dataset(latest_nabu_node[path])

    def cancel(self):
        """
        stop current processing
        """
        self._cancelled = True
        if self._current_processing is not None:
            self._current_processing.cancel()


class NabuVolume(NabuVolumeTask):
    def __init__(
        self,
        process_id=None,
        varinfo=None,
        inputs=None,
        node_id=None,
        node_attrs=None,
        execinfo=None,
    ):
        deprecated_warning(
            name="tomwer.core.process.reconstruction.nabu.nabuvolume.NabuVolume",
            type_="class",
            reason="improve readibility",
            since_version="1.2",
            replacement="NabuVolumeTask",
        )
        super().__init__(process_id, varinfo, inputs, node_id, node_attrs, execinfo)
