import os

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pint
from ewokscore import Task

# from PIL import Image
from silx.io import h5py_utils
from silx.io import url
from tifffile import RESUNIT
from tifffile import imwrite

ureg = pint.UnitRegistry()

DEFAULTS = {
    "do_gallery": False,
    "output_template": "IMG_{elm}_areal_density_ng_mm2_{scanno:04d}",
}


class TiffExporterFromRegrid(
    Task,
    input_names=[
        "bliss_scan_uri",
        "output_root_uri",
        "xrf_results_uri",
        "output_root_group",
    ],
    optional_input_names=[
        "do_gallery",
        "output_template",
    ],
    output_names=[
        "bliss_scan_uri",
        "output_root_uri",
        "xrf_results_uri",
        "output_root_group",
    ],
):
    def run(self):

        params = {**DEFAULTS, **self.get_input_values()}

        self.outputs.bliss_scan_uri = params["bliss_scan_uri"]
        self.outputs.output_root_uri = params["output_root_uri"]
        self.outputs.xrf_results_uri = params["xrf_results_uri"]
        self.outputs.output_root_group = params["output_root_group"]

        do_gallery = params["do_gallery"]
        output_template = params["output_template"]
        output_root_group = params["output_root_group"]

        xrf_results_uri = url.DataUrl(params["xrf_results_uri"])

        output_path = os.path.dirname(xrf_results_uri.file_path())
        output_gallery_path = os.path.join(output_path, "gallery")

        os.makedirs(output_path, exist_ok=True)
        os.chmod(
            output_path, 0o770
        )  # nosec B103: group write required for shared pipeline

        if do_gallery:
            os.makedirs(output_gallery_path, exist_ok=True)
            os.chmod(
                output_gallery_path, 0o770
            )  # nosec B103: group write required for shared pipeline

        bliss_scan_uri = url.DataUrl(self.inputs.bliss_scan_uri)

        uri = bliss_scan_uri.data_path()
        uri = uri.split("/")[1]
        scanno = int(uri.split(".")[0])

        print(xrf_results_uri)

        with h5py_utils.open_item(xrf_results_uri.file_path(), xrf_results_uri.data_path()) as _res:  # type: ignore[reportGeneralTypeIssues]

            quanti = False

            if "massfractions" in _res:
                _res = _res["massfractions"]
                quanti = True
            elif "parameters" in _res:
                _res = _res["parameters"]
            else:
                raise RuntimeError(
                    "Found none of massfractions or parameters to be exported!"
                )

            grps = []
            if "signal" in _res.attrs:
                grps += [
                    _res.attrs["signal"],
                ]

            if "auxiliary_signals" in _res.attrs:
                grps += list(_res.attrs["auxiliary_signals"])

            kwargs = {}
            axes = []
            axnames = []
            if "axes" in _res.attrs:
                resolutions = []
                resolutions_units = []

                for a in _res.attrs["axes"]:
                    axnames += [
                        _res[a].attrs["long_name"],
                    ]
                    d = np.array(_res[a])
                    axes += [
                        d,
                    ]

                    res = float(np.mean(np.diff(d)))
                    resu = _res[a].attrs["units"]

                    resolutions_units += [
                        resu,
                    ]

                    ps = ureg.Quantity(res, resu)

                    if ps.check("[length]"):
                        resolutions += [
                            1.0 / ps.to("cm").magnitude,
                        ]
                    else:
                        resolutions += [
                            1e4 / ps.magnitude,
                        ]

                kwargs["resolutionunit"] = RESUNIT.CENTIMETER
                kwargs["resolution"] = tuple(resolutions)
                axes = axes[::-1]
                axnames = axnames[::-1]

            cmap = "jet"

            for g in grps:

                data = np.array(_res[g], dtype=np.float32)
                #                im = Image.fromarray(data, mode="F")
                gg = g.replace(" ", "_")
                templatedata = {
                    "elm": gg,
                    "detector_name": output_root_group,
                    "scanno": scanno,
                }
                #                im.save(
                #                    os.path.join(
                #                        output_path, output_template.format(**templatedata) + ".tiff"
                #                    ),
                #                    **kwargs,
                #                )

                imwrite(
                    os.path.join(
                        output_path, output_template.format(**templatedata) + ".tiff"
                    ),
                    data,
                    **kwargs,
                )

                if output_gallery_path is not None:
                    norm = mpl.colors.Normalize(0, np.nanmax(data))
                    f, ax = plt.subplots()
                    ax.pcolor(*axes, data, norm=norm, cmap=cmap, shading="nearest")
                    ax.invert_yaxis()

                    if "deg" not in resolutions_units:
                        ax.set_aspect("equal", "box")

                    if len(axnames) == 2:
                        ax.set_xlabel(axnames[0])
                        ax.set_ylabel(axnames[1])

                    f.suptitle(g)
                    f.colorbar(
                        mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
                        ax=ax,
                        label=(
                            "Areal mass density ($ng/mm^2$)" if quanti else "Parameters"
                        ),
                    )
                    f.savefig(
                        os.path.join(
                            output_gallery_path,
                            output_template.format(**templatedata) + ".jpg",
                        )
                    )
