import numpy as np

import czifile as czi
import xml.etree.ElementTree as ET


class ZeissData:

    def __init__(self, filename):

        self.filename = filename
        self._initialize()

    def _initialize(self):
        czi_obj = czi.CziFile(self.filename)
        axes = czi_obj.axes
        tmp = czi_obj.asarray()
        self.axes = ""
        # print(tmp.shape)
        for i, size in enumerate(tmp.shape):
            if size > 1:
                self.axes += axes[i]

        self.image = np.squeeze(tmp)
        self.shape = self.image.shape
        self.ndim = len(self.shape)
        self.origin = np.array([[0, i1] for i1 in self.shape])

        self.metadata = czi_obj.metadata()
        root = ET.fromstring(self.metadata)
        for d in root.findall("./Metadata/Scaling/Items/Distance"):
            el = d.attrib["Id"]
            if el in self.axes:
                setattr(self, "p" + el.lower(), float(d[0].text) * 1e6)

        # print(self.axes)
        # print(self.shape)

        # for k, v in kwargs.items():
        #     if k in ImageJ.allowed_kwargs:
        #         setattr(self, k, v)
        #     else:
        #         print(k, "is not allowed keyword argument!")

    def get_channle_info(self, channelID):
        #  print(self.metadata)
        root = ET.fromstring(self.metadata)
        # for child in root:
        #     print(child.tag, child.attrib)
        # print(type(root))

        na = None
        refractive_index = None
        for _child in root.findall("./Metadata/Information/Instrument/Objectives/Objective"):
            for _ch in _child:
                if _ch.tag == "LensNA":
                    na = _ch.text

                if _ch.tag == "ImmersionRefractiveIndex":
                    refractive_index = _ch.text
            # print(_child.tag, _child.attrib, _child[1].tag, type(_child.tag), type(_child.attrib))

        em, ex = None, None
        for _child in root.findall("./Metadata/DisplaySetting/Channels"):

            for _ch in _child.findall("./Channel"):
                if _ch.attrib["Id"] == f"Channel:{channelID}":
                    for _c in _ch:
                        if _c.tag == "DyeMaxEmission":
                            em = _c.text
                        if _c.tag == "DyeMaxExcitation":
                            ex = _c.text
                        if ex is not None and em is not None:
                            break

            #             print(_c.tag)
            #         print(_ch.tag, _ch.attrib["Id"], _ch.text)

            # print()
            # print(_child.tag, _child.attrib, _child[1].tag, type(_child.tag), type(_child.attrib))

        info = {
            "MainEmissionWavelength": float(em),
            "MainExcitationWavelength": float(ex),
            "ObjectiveNA": float(na),
            "RefractiveIndex": float(refractive_index),
        }
        return info

    def get_stack(self, channleID):
        #spacings = dict(X=self.px * 1e-6, Y=self.py * 1e-6, Z=self.pz * 1e-6)
        spacings = dict(X=self.px * 1e-6, Y=self.py * 1e-6)

        return self.image[channleID], spacings, self.get_channle_info(channleID)