# -*- coding: utf-8 -*-
"""Module to plot feature output from forward passes for visual inspection"""
import glob
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime as dt

import matplotlib.pyplot as plt
import numpy as np
import rex
from rex.utilities.fun_utils import get_fun_call_str

from sup3r.utilities import ModuleName
from sup3r.utilities.cli import BaseCLI

logger = logging.getLogger(__name__)


class Sup3rVisualQa:
    """Module to plot features for visual qa"""

    def __init__(
        self,
        file_paths,
        out_pattern,
        features,
        time_step=10,
        spatial_slice=slice(None),
        source_handler_class=None,
        max_workers=None,
        overwrite=False,
        **kwargs,
    ):
        """
        Parameters
        ----------
        file_paths : list | str
            Specifies the files to use for the plotting routine. This is either
            a list of h5 files generated by the forward pass module or a string
            pointing to h5 forward pass output which can be parsed by glob.glob
        out_pattern : str
            The pattern to use for naming the plot figures. This must include
            {feature} and {index} so output files can be named with
            out_pattern.format(feature=feature, index=index).
            e.g. outfile_{feature}_{index}.png. The number of plot figures is
            determined by the time_index of the h5 files and the time_step
            argument. The index key refers to the plot file index from the list
            of all plot files generated.
        features : list
            List of features to plot from the h5 files provided.
        time_step : int
            Number of timesteps to average over for a single plot figure.
        spatial_slice : slice
            Slice specifying the spatial range to plot. This can include a
            step > 1 to speed up plotting.
        source_handler_class : str | None
            Name of the class to use for h5 input files. If None this defaults
            to MultiFileResource.
        max_workers : int | None
            Max number of workers to use for plotting. If workers=1 then all
            plots will be created in serial.
        overwrite : bool
            Whether to overwrite saved plots.
        **kwargs : dict
            Dictionary of kwargs passed to matplotlib.pyplot.scatter().
        """

        self.features = features
        self.out_pattern = out_pattern
        self.time_step = time_step
        self.spatial_slice = (
            spatial_slice
            if isinstance(spatial_slice, slice)
            else slice(*spatial_slice)
        )
        self.file_paths = (
            file_paths
            if isinstance(file_paths, list)
            else glob.glob(file_paths)
        )
        self.max_workers = max_workers
        self.kwargs = kwargs
        self.res_handler = source_handler_class or 'MultiFileResource'
        self.res_handler = getattr(rex, self.res_handler)
        self.overwrite = overwrite
        if not os.path.exists(os.path.dirname(out_pattern)):
            os.makedirs(os.path.dirname(out_pattern), exist_ok=True)
        logger.info(
            'Initializing Sup3rVisualQa with '
            f'file_paths={self.file_paths}, '
            f'out_pattern={self.out_pattern}, '
            f'features={self.features}, '
            f'time_step={self.time_step}, '
            f'spatial_slice={self.spatial_slice}, '
            f'source_handler_class={self.res_handler}, '
            f'max_workers={max_workers}, '
            f'overwrite={self.overwrite}, '
            f'kwargs={kwargs}.'
        )

    def run(self):
        """
        Create plot figures for all the features in self.features. For each
        feature there will be n_files created, where n_files is the number of
        timesteps in the h5 files provided divided by self.time_step.
        """
        with self.res_handler(self.file_paths) as res:
            time_index = res.time_index
            n_files = len(time_index[:: self.time_step])
            time_slices = np.array_split(np.arange(len(time_index)), n_files)
            time_slices = [slice(s[0], s[-1] + 1) for s in time_slices]

            if self.max_workers == 1:
                self._serial_figure_plots(
                    res, time_index, time_slices, self.spatial_slice
                )
            else:
                self._parallel_figure_plots(
                    res, time_index, time_slices, self.spatial_slice
                )

    def _serial_figure_plots(
        self, res, time_index, time_slices, spatial_slice
    ):
        """Plot figures in parallel with max_workers=self.workers

        Parameters
        ----------
        res : MultiFileResourceX
            Resource handler for the provided h5 files
        time_index : pd.DateTimeIndex
            The time index for the provided h5 files
        time_slices : list
             List of slices specifying all the time ranges to average and plot
        spatial_slice : slice
             Slice specifying the spatial range to plot
        """
        for feature in self.features:
            for i, t_slice in enumerate(time_slices):
                out_file = self.out_pattern.format(
                    feature=feature, index=str(i).zfill(8)
                )
                self.plot_figure(
                    res, time_index, feature, t_slice, spatial_slice, out_file
                )

    def _parallel_figure_plots(
        self, res, time_index, time_slices, spatial_slice
    ):
        """Plot figures in parallel with max_workers=self.workers

        Parameters
        ----------
        res : MultiFileResourceX
            Resource handler for the provided h5 files
        time_index : pd.DateTimeIndex
            The time index for the provided h5 files
        time_slices : list
             List of slices specifying all the time ranges to average and plot
        spatial_slice : slice
             Slice specifying the spatial range to plot
        """
        futures = {}
        now = dt.now()
        n_files = len(time_slices) * len(self.features)
        with ThreadPoolExecutor(max_workers=self.max_workers) as exe:
            for feature in self.features:
                for i, t_slice in enumerate(time_slices):
                    out_file = self.out_pattern.format(
                        feature=feature, index=str(i).zfill(8)
                    )
                    future = exe.submit(
                        self.plot_figure,
                        res,
                        time_index,
                        feature,
                        t_slice,
                        spatial_slice,
                        out_file,
                    )
                    futures[future] = out_file

            logger.info(
                f'Started plotting {n_files} files ' f'in {dt.now() - now}.'
            )

            for i, future in enumerate(as_completed(futures)):
                try:
                    future.result()
                except Exception as e:
                    msg = f'Error making plot {futures[future]}.'
                    logger.exception(msg)
                    raise RuntimeError(msg) from e
                logger.debug(f'{i+1} out of {n_files} plots created.')

    def plot_figure(
        self, res, time_index, feature, t_slice, s_slice, out_file
    ):
        """Plot temporal average for the given feature and with the time range
        specified by t_slice

        Parameters
        ----------
        res : MultiFileResourceX
            Resource handler for the provided h5 files
        time_index : pd.DateTimeIndex
            The time index for the provided h5 files
        feature : str
            The feature to plot
        t_slice : slice
            The slice specifying the time range to average and plot
        s_slice : slice
            The slice specifying the spatial range to plot.
        out_file : str
            Name of the output plot file
        """
        if not self.overwrite and os.path.exists(out_file):
            logger.info(
                f'{out_file} already exists and overwrite='
                f'{self.overwrite}. Skipping this plot.'
            )
            return
        start_time = time_index[t_slice.start]
        stop_time = time_index[t_slice.stop - 1]
        logger.info(
            f'Plotting time average for {feature} from '
            f'{start_time} to {stop_time}.'
        )
        fig = plt.figure()
        title = f'{feature}: {start_time} - {stop_time}'
        plt.suptitle(title)
        plt.scatter(
            res.meta.longitude[s_slice],
            res.meta.latitude[s_slice],
            c=np.mean(res[feature, t_slice, s_slice], axis=0),
            **self.kwargs,
        )
        plt.colorbar()
        fig.savefig(out_file)
        plt.close()
        logger.info(f'Saved figure {out_file}')

    @classmethod
    def get_node_cmd(cls, config):
        """Get a CLI call to initialize Sup3rVisualQa and execute the
        Sup3rVisualQa.run() method based on an input config

        Parameters
        ----------
        config : dict
            sup3r QA config with all necessary args and kwargs to
            initialize Sup3rVisualQa and execute Sup3rVisualQa.run()
        """
        import_str = 'import time;\n'
        import_str += 'from gaps import Status;\n'
        import_str += 'from rex import init_logger;\n'
        import_str += 'from sup3r.qa.visual_qa import Sup3rVisualQa;\n'

        qa_init_str = get_fun_call_str(cls, config)

        log_file = config.get('log_file', None)
        log_level = config.get('log_level', 'INFO')

        log_arg_str = f'"sup3r", log_level="{log_level}"'
        if log_file is not None:
            log_arg_str += f', log_file="{log_file}"'

        cmd = (
            f"python -c \'{import_str}\n"
            "t0 = time.time();\n"
            f"logger = init_logger({log_arg_str});\n"
            f"qa = {qa_init_str};\n"
            "qa.run();\n"
            "t_elap = time.time() - t0;\n"
        )

        pipeline_step = config.get('pipeline_step') or ModuleName.VISUAL_QA
        cmd = BaseCLI.add_status_cmd(config, pipeline_step, cmd)
        cmd += ";\'\n"

        return cmd.replace('\\', '/')
