from ewokscore import Task
import logging
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg
from pathlib import Path
from ewoksid02.id02_format import open_id02_file

logger = logging.getLogger(__name__)


class SaveGalleryTask(
    Task,
    input_names=[
        "nxdata_url",
    ],
    optional_input_names=[
        "do_save",
    ],
    output_names=[
        "filename_png",
    ],
):
    def run(self):
        if not self.get_input_value("do_save", True):
            logger.info("Saving in gallery is skipped")
            return

        nxdata_url = self.inputs.nxdata_url
        filename_processed, h5path = nxdata_url.split("::")
        processed_directory = Path(filename_processed).parent
        gallery_directory = processed_directory / "gallery"
        gallery_directory.mkdir(exist_ok=True)
        filename_png = (
            gallery_directory / Path(filename_processed).with_suffix(".png").name
        )
        filename_png = str(filename_png)

        with open_id02_file(filename=filename_processed) as file:
            if h5path not in file:
                logger.error(f"No group {h5path} in {filename_processed}")
                return
            nxdata_group = file[h5path]
            if "data" not in nxdata_group:
                logger.error(f"No data dataset in {filename_processed}::{h5path}")
                return
            dset = nxdata_group["data"]
            nb_frames = len(dset)
            data_mean = dset[:].mean(axis=0)
        fig = Figure()
        FigureCanvasAgg(fig)
        ax = fig.subplots()
        try:
            ax.set_title(f"{str(Path(filename_processed).stem)} ({nb_frames} frames)")
            ax.set_xlabel("q ($nm^{-1}$)")
            ax.set_ylabel("Intensity")
            if data_mean.ndim == 1:
                ax.loglog(data_mean)
            elif data_mean.ndim == 2:
                ax.imshow(data_mean)
            fig.savefig(filename_png)
        finally:
            fig.clear()
        self.outputs.filename_png = filename_png

        logger.info(f"Saving in the gallery: {filename_png}\n")
