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
        print("is file", self.roifile_fn.is_file())
        self.rois = ImagejRoi.fromfile(self.roifile_fn)
        self.stem = self.roifile_fn.stem.split("Roi-")[1]
        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.lower()}*.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]))

        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,
            }
        print("Image props", self.image_props)

        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].split("p")[1])
                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 i, roi in enumerate(self.rois):
                rc = roi.subpixel_coordinates
                if rc is None:
                    print(
                        f"No coordinates: {self.roifile_fn}, element {i}, out of {len(self.rois)}"
                    )
                    continue
                xmn, ymn = rc.min(axis=0)
                xmx, ymx = rc.max(axis=0)
                xc = (xmn + xmx) / 2
                yc = (ymn + ymx) / 2
                bw = xmx - xmn
                bh = ymx - ymn
                coords = f"{xc/self.width} {yc/self.height} {bw/self.width} {bh/self.height} "
                for x, y in rc:
                    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("-i", "--input", nargs="*", type=Path)
    parser.add_argument("-o", "--output", type=Path)
    args = parser.parse_args()

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

        print()