import numpy as np

from roifile import ImagejRoi
from tifffile import TiffFile, TiffWriter
from pathlib import Path


class UT:
    """
    Docstring for UT

    Operetta files along with rois drawn in ImageJ
    """

    def __init__(self, roifile_fn: Path):
        self.roifile_fn = roifile_fn
        self.rois = ImagejRoi.fromfile(self.roifile_fn)
        self.stem = self.roifile_fn.stem.strip("-RoiSet")
        self.image, self.image_props = self._load_images()

    def _load_images(self):
        """Loading sequence of tif files
        array sequence is CZYX
        """
        print(self.roifile_fn.parent, self.stem)
        fns = list(self.roifile_fn.parent.glob(f"{self.stem}*.tif*"))
        stems = [fn.stem.split(self.stem)[-1] for fn in fns]
        n_ch = len(set([stem.split("-ch")[-1].split("t")[0] for stem in stems]))
        n_p = len(set([stem.split("-")[0] for stem in stems]))
        n_t = len(set([stem.split("t")[1] for stem in stems]))
        print(n_ch, n_p, n_t)

        with TiffFile(fns[0]) as tif:
            img = tif.asarray()
            w, h = img.shape
            dtype = img.dtype
            self.image_props = {
                "channels": n_ch,
                "planes": n_p,
                "tiles": n_t,
                "width": w,
                "height": h,
                "dtype": dtype,
            }

        image_stack = np.zeros((n_ch, n_p, w, h), dtype=dtype)
        for fn in fns:
            with TiffFile(fn) as tif:
                img = tif.asarray()
                stem = fn.stem.split(self.stem)[-1]
                ch = int(stem.split("-ch")[-1].split("t")[0])
                p = int(stem.split("-")[0].lstrip("p"))
                t = int(stem.split("t")[1])
                print(fn.stem, "ch", ch, "p", p, "t", t)
                image_stack[ch - 1, p - 1] = img

        print(image_stack.shape)

        return image_stack, self.image_props

    @property
    def width(self):
        return self.image_props["width"]

    @property
    def height(self):
        return self.image_props["height"]

    @property
    def nchannels(self):
        return self.image_props["channels"]

    @property
    def nplanes(self):
        return self.image_props["planes"]

    def export_rois(
        self,
        path: Path,
        subfolder: str = "labels",
        class_index: int = 0,
    ):
        """Export rois to a file"""
        with open(path / subfolder / f"{self.stem}.txt", "w") as f:
            for roi in self.rois:
                # TODO add image coordinates normalization
                coords = ""
                for x, y in roi.subpixel_coordinates:
                    coords += f"{x/self.width} {y/self.height}"
                f.write(f"{class_index} {coords}\n")

        return

    def export_image(
        self,
        path: Path,
        subfolder: str = "images",
        plane_mode: str = "max projection",
        channel: int = 0,
    ):
        """Export image to a file"""

        if plane_mode == "max projection":
            self.image = np.max(self.image[channel], axis=0)
            print(self.image.shape)

        with TiffWriter(path / subfolder / f"{self.stem}.tif") as tif:
            tif.write(self.image)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("input", type=Path)
    parser.add_argument("output", type=Path)
    args = parser.parse_args()

    for rfn in args.input.glob("*.zip"):
        ut = UT(rfn)
        ut.export_rois(args.output, class_index=0)
        ut.export_image(args.output, plane_mode="max projection", channel=0)