import numpy as np
from construct import (
    PaddedString,
    Int16un,
    Struct,
    Int8un,
    Int32sn,
    Int32un,
    Int64un,
    Array,
    Float32l,
)
from oct_converter.image_types import OCTVolumeWithMetaData, FundusImageWithMetaData
#import julian
import datetime
import collections


class E2E(object):
    """Class for extracting data from Heidelberg's .e2e file format.

    Notes:
        Mostly based on description of .e2e file format here:
        https://bitbucket.org/uocte/uocte/wiki/Heidelberg%20File%20Format.

    Attributes:
        filepath (str): Path to .img file for reading.
        header_structure (obj:Struct): Defines structure of volume's header.
        main_directory_structure (obj:Struct): Defines structure of volume's main directory.
        sub_directory_structure (obj:Struct): Defines structure of each sub directory in the volume.
        chunk_structure (obj:Struct): Defines structure of each data chunk.
        image_structure (obj:Struct): Defines structure of image header.
    """

    e2e_metadata_dict = {}

    def __init__(self, filepath):
        self.filepath = filepath
        self.header_structure = Struct(
            "magic" / PaddedString(12, "ascii"),
            "version" / Int32un,
            "unknown" / Array(10, Int16un),
        )
        self.main_directory_structure = Struct(
            "magic" / PaddedString(12, "ascii"),
            "version" / Int32un,
            "unknown" / Array(10, Int16un),
            "num_entries" / Int32un,
            "current" / Int32un,
            "prev" / Int32un,
            "unknown3" / Int32un,
        )
        self.sub_directory_structure = Struct(
            "pos" / Int32un,
            "start" / Int32un,
            "size" / Int32un,
            "unknown" / Int32un,
            "patient_id" / Int32un,
            "study_id" / Int32un,
            "series_id" / Int32un,
            "slice_id" / Int32sn,
            "unknown2" / Int16un,
            "unknown3" / Int16un,
            "type" / Int32un,
            "unknown4" / Int32un,
        )
        self.chunk_structure = Struct(
            "magic" / PaddedString(12, "ascii"),
            "unknown" / Int32un,
            "unknown2" / Int32un,
            "pos" / Int32un,
            "size" / Int32un,
            "unknown3" / Int32un,
            "patient_id" / Int32un,
            "study_id" / Int32un,
            "series_id" / Int32un,
            "slice_id" / Int32sn,
            "ind" / Int16un,
            "unknown4" / Int16un,
            "type" / Int32un,
            "unknown5" / Int32un,
        )
        self.image_structure = Struct(
            "size" / Int32un,
            "type" / Int32un,
            "unknown" / Int32un,
            "width" / Int32un,
            "height" / Int32un,
        )
        """
        self.patient_struct = Struct(
            "name" / PaddedString(31, "ascii"),
            "surname" / PaddedString(66, "ascii"),
            "dob" / Int32un,
            "sex" / PaddedString(1, "ascii"),
            "patient_id" / Int32un,
        )
        """
        self.patient_struct = Struct(
            "first_name" / PaddedString(31, "ascii"),
            "surname" / PaddedString(51, "ascii"),
            "title" / PaddedString(15, "ascii"),
            "birthdate" / Int32un,
            "sex" / PaddedString(1, "ascii"),
            "patient_id" / PaddedString(25, "ascii"),
        )
            
        self.slodataelement_struct = Struct(
            "datetime" / Int64un,
            "transform" / Array(6, Float32l),
        )
        self.lat_structure = Struct(
            "unknown" / Array(14, Int8un),
            "laterality" / Int8un,
            "unknown2" / Int8un,
        )
        self.bscan_metadata_struct = Struct(
            "unknown1" / Int32un,
            "imgSizeX" / Int32un,
            "imgSizeY" / Int32un,
            "posX1" / Float32l,
            "posY1" / Float32l,
            "posX2" / Float32l,
            "posY2" / Float32l,
            "zero1" / Int32un,
            "unknown2" / Float32l,
            "scaleY" / Float32l,
            "unknown3" / Float32l,
            "zero2" / Int32un,
            "unknown4" / Array(2, Float32l),
            "zero3" / Int32un,
            "imgSizeWidth" / Int32un,
            "numImages" / Int32un,
            "aktImage" / Int32un,
            "scantype" / Int32un,
            "centerPosX" / Float32l,
            "centerPosY" / Float32l,
            "unknown5" / Int32un,
            "acquisitionTime" / Int64un,
            "unknown6" / Int32un,
            "numAve" / Int32un,
            "unknown7" / Int32un,
            "imageQuality" / Float32l,
        )

    def write_deidentified(self, outfile_path):
        # read file to get all patient_struct locations
        self.read_oct_volume()

        if len(self.patient_struct_locs) != 1:
            print(f"ERROR WITH FILE {self.filepath}")
            print(f"patient struct locations: {self.patient_struct_locs}")
            return

        with open(self.filepath, "rb") as f:
            with open(outfile_path, "wb") as out_f:
                # iterate all bscan meta data location
                end = 0
                for bscan_struct_loc in sorted(self.bscan_struct_locs):
                    out_f.write(f.read(bscan_struct_loc - end))
                    raw = f.read(self.bscan_metadata_struct.sizeof())
                    bscan_data = self.bscan_metadata_struct.parse(raw)
                    
                    # use the original value except acquisitionTime
                    unknown1 = self.bscan_metadata_struct._subcons["unknown1"].build(bscan_data.unknown1)
                    imgSizeX = self.bscan_metadata_struct._subcons["imgSizeX"].build(bscan_data.imgSizeX)
                    imgSizeY = self.bscan_metadata_struct._subcons["imgSizeY"].build(bscan_data.imgSizeY)
                    posX1 = self.bscan_metadata_struct._subcons["posX1"].build(bscan_data.posX1)
                    posY1 = self.bscan_metadata_struct._subcons["posY1"].build(bscan_data.posY1)
                    posX2 = self.bscan_metadata_struct._subcons["posX2"].build(bscan_data.posX2)
                    posY2 = self.bscan_metadata_struct._subcons["posY2"].build(bscan_data.posY2)
                    zero1 = self.bscan_metadata_struct._subcons["zero1"].build(bscan_data.zero1)
                    unknown2 = self.bscan_metadata_struct._subcons["unknown2"].build(bscan_data.unknown2)
                    scaleY = self.bscan_metadata_struct._subcons["scaleY"].build(bscan_data.scaleY)
                    unknown3 = self.bscan_metadata_struct._subcons["unknown3"].build(bscan_data.unknown3)
                    zero2 = self.bscan_metadata_struct._subcons["zero2"].build(bscan_data.zero2)
                    unknown4 = self.bscan_metadata_struct._subcons["unknown4"].build(bscan_data.unknown4)
                    zero3 = self.bscan_metadata_struct._subcons["zero3"].build(bscan_data.zero3)
                    imgSizeWidth = self.bscan_metadata_struct._subcons["imgSizeWidth"].build(bscan_data.imgSizeWidth)
                    numImages = self.bscan_metadata_struct._subcons["numImages"].build(bscan_data.numImages)
                    aktImage = self.bscan_metadata_struct._subcons["aktImage"].build(bscan_data.aktImage)
                    scantype = self.bscan_metadata_struct._subcons["scantype"].build(bscan_data.scantype)
                    centerPosX = self.bscan_metadata_struct._subcons["centerPosX"].build(bscan_data.centerPosX)
                    centerPosY = self.bscan_metadata_struct._subcons["centerPosY"].build(bscan_data.centerPosY)
                    unknown5 = self.bscan_metadata_struct._subcons["unknown5"].build(bscan_data.unknown5)
                    acquisitionTime = self.bscan_metadata_struct._subcons["acquisitionTime"].build(116444720880000000)
                    unknown6 = self.bscan_metadata_struct._subcons["unknown6"].build(bscan_data.unknown6)
                    numAve = self.bscan_metadata_struct._subcons["numAve"].build(bscan_data.numAve)
                    unknown7 = self.bscan_metadata_struct._subcons["unknown7"].build(bscan_data.unknown7)
                    imageQuality = self.bscan_metadata_struct._subcons["imageQuality"].build(bscan_data.imageQuality)

                    items = [unknown1, imgSizeX, imgSizeY, posX1, posY1, posX2, posY2, zero1, unknown2, scaleY, unknown3,
                            zero2, unknown4, zero3, imgSizeWidth, numImages, aktImage, scantype, centerPosX, centerPosY,
                            unknown5, acquisitionTime, unknown6, numAve, unknown7, imageQuality]
                    for item in items:
                        out_f.write(item)
                        
                    end = bscan_struct_loc + self.bscan_metadata_struct.sizeof() 

                # end of for loop
                
                patient_struct_loc = self.patient_struct_locs[0]
                if len(self.bscan_struct_locs)==0:
                    return "Can not find B-Scan..."
                if patient_struct_loc < sorted(self.bscan_struct_locs)[-1]:
                    return "Patient Info comes earlier than B-Scan..."
                
                # write up until patient struct
                out_f.write(f.read(patient_struct_loc - end))
                # f.seek(patient_struct_loc)

                raw = f.read(self.patient_struct.sizeof())
                pt_data = self.patient_struct.parse(raw)

                first_name = self.patient_struct._subcons["first_name"].build("FIRSTNAME")
                last_name = self.patient_struct._subcons["surname"].build("LASTNAME")
                title = self.patient_struct._subcons["title"].build("TITLE")
                dob = self.patient_struct._subcons["birthdate"].build(19000101)
                sex = self.patient_struct._subcons["sex"].build("N")
                patient_id = self.patient_struct._subcons["patient_id"].build("PATIENT_ID")
                

                out_f.write(first_name)
                out_f.write(last_name)
                out_f.write(title)
                out_f.write(dob)
                out_f.write(sex)
                out_f.write(patient_id)
                # write rest of file
                out_f.write(f.read())
                return 0 

    def get_metadata(self, f):
        MAGIC_OFFSET = 116444720880000000  # 369 years
        raw = f.read(self.bscan_metadata_struct.sizeof())
        bscan_metadata = self.bscan_metadata_struct.parse(raw)
        timestamp = bscan_metadata.acquisitionTime
        timestamp_sec = int((timestamp - MAGIC_OFFSET) / 10000000)

        converted_date = datetime.datetime.fromtimestamp(timestamp_sec)

        return {
            "timestamp": converted_date,
            "image_index": bscan_metadata.aktImage,
            "num_slices": bscan_metadata.numImages,
            "scantype": bscan_metadata.scantype,
            "width": bscan_metadata.imgSizeX,
            "height": bscan_metadata.imgSizeY,
            "scaleY": bscan_metadata.scaleY,
            "x1": bscan_metadata.posX1,
            "y1": bscan_metadata.posY1,
            "x2": bscan_metadata.posX2,
            "y2": bscan_metadata.posY2,
        }

    def read_oct_volume(self, uint8_out=True):
        """Reads OCT data.

        Returns:
            obj:OCTVolumeWithMetaData
        """
        DATE_INT_OFFSET = 14555840.5

        most_recent_bscan_metadata_dict = {}
        oct_volume_dict = {}
        patient_metadata_dict = {}
        header_offset = 0
        self.patient_struct_locs = []
        self.bscan_struct_locs = []
        self.laterality = None

        with open(self.filepath, "rb") as f:
            try:
                raw = f.read(self.header_structure.sizeof())
                header = self.header_structure.parse(raw)

                raw = f.read(self.main_directory_structure.sizeof())
                main_directory = self.main_directory_structure.parse(raw)
            except UnicodeDecodeError:
                header_offset = 0x40
                f.seek(header_offset)
                raw = f.read(self.header_structure.sizeof())
                header = self.header_structure.parse(raw)

                raw = f.read(self.main_directory_structure.sizeof())
                main_directory = self.main_directory_structure.parse(raw)

            # traverse list of main directories in first pass
            directory_stack = []

            current = main_directory.current
            while current != 0:
                directory_stack.append(current)
                f.seek(current + header_offset)
                raw = f.read(self.main_directory_structure.sizeof())
                directory_chunk = self.main_directory_structure.parse(raw)
                current = directory_chunk.prev

            # traverse in second pass and  get all subdirectories
            chunk_stack = []
            volume_dict = {}
            for position in directory_stack:
                f.seek(position + header_offset)
                raw = f.read(self.main_directory_structure.sizeof())
                directory_chunk = self.main_directory_structure.parse(raw)

                for ii in range(directory_chunk.num_entries):
                    raw = f.read(self.sub_directory_structure.sizeof())
                    chunk = self.sub_directory_structure.parse(raw)

                    if chunk.type == 0x40000000:
                        volume_string = "{}_{}_{}".format(
                            chunk.patient_id, chunk.study_id, chunk.series_id
                        )
                        if volume_string not in volume_dict.keys():
                            volume_dict[volume_string] = chunk.slice_id / 2
                        elif chunk.slice_id / 2 > volume_dict[volume_string]:
                            volume_dict[volume_string] = chunk.slice_id / 2

                    if chunk.start > chunk.pos:
                        chunk_stack.append([chunk.start, chunk.size])

            # initalise dict to hold all the image volumes
            volume_array_dict = {}
            for volume, num_slices in volume_dict.items():
                if num_slices > 0:
                    volume_array_dict[volume] = [0] * int(num_slices)

            for start, pos in chunk_stack:
                f.seek(start + header_offset)
                raw = f.read(self.chunk_structure.sizeof())
                try:
                    chunk = self.chunk_structure.parse(raw)
                except Exception as e:
                    continue

                if chunk.type == 0x2714:  # bscan metadata
                    most_recent_bscan_metadata_dict = self.get_metadata(f)
                    bscan_struct_loc = start + 60 + header_offset
                    self.bscan_struct_locs.append(bscan_struct_loc)

                if chunk.type == 0x00000009:
                    patient_struct_loc = start + 60 + header_offset
                    self.patient_struct_locs.append(patient_struct_loc)
                    f.seek(patient_struct_loc)
                    raw = f.read(self.patient_struct.sizeof())
                    pt_data = self.patient_struct.parse(raw)
                    #try:
                    #    dob = julian.from_jd((pt_data.dob / 64.0) - DATE_INT_OFFSET)
                    #except Exception as e:
                    dob = None
                    patient_metadata_dict.update(
                        {
                            "dob": dob,
                            "name": pt_data.first_name,
                            "surname": pt_data.surname,
                            "sex": pt_data.sex,
                        }
                    )

                if chunk.type == 0xB:  # laterality data
                    raw = f.read(self.lat_structure.sizeof())
                    try:
                        laterality_data = self.lat_structure.parse(raw)
                        if laterality_data.laterality == 82:
                            self.laterality = "R"
                        elif laterality_data.laterality == 76:
                            self.laterality = "L"
                    except:
                        self.laterality = None

                if chunk.type == 0x40000000:  # image data
                    raw = f.read(self.image_structure.sizeof())
                    image_data = self.image_structure.parse(raw)

                    volume_string = "{}_{}_{}".format(
                        chunk.patient_id, chunk.study_id, chunk.series_id
                    )
                    slice_index = int(chunk.slice_id / 2) - 1

                    if chunk.ind == 1:  # oct data
                        # read entire image chunk
                        byte_chunk = f.read(2 * image_data.height * image_data.width)

                        # read bytes
                        image_byte_array = np.frombuffer(byte_chunk, dtype=np.uint16)

                        # reverse bits and repack
                        image_byte_array = np.packbits(
                            np.unpackbits(image_byte_array.view(np.uint8))[::-1]
                        ).view(np.uint16)[::-1]

                        # grab bottom 6 bits and reverse order
                        exponent_array = image_byte_array & 0x3F
                        # exponent_array = np.packbits(np.unpackbits(exponent_array.view(np.uint8))[::-1]).view(np.uint16)[::-1] >> 10
                        exponent_array = (
                            np.packbits(
                                np.unpackbits(exponent_array.view(np.uint8))[::-1]
                            ).view(np.uint16)[::-1]
                            >> 10
                        )
                        # manipulate exponent
                        exponent_array = exponent_array.view(np.int16) - 63

                        # grab the top 10 bits & manipulate mantissa
                        mantissa_array = 1 + (image_byte_array >> 6) / 1024

                        unshaped_image = mantissa_array * np.float_power(
                            2, exponent_array
                        )
                        # image = 256 * pow(image, 1.0 / 2.4)
                        unshaped_image = np.float_power(unshaped_image, 1.0 / 3.8)

                        # reshape image into width and height size
                        try:
                            image = unshaped_image.reshape(
                                image_data.width, image_data.height
                            )
                        except Exception as e:
                            break

                        if uint8_out:
                            # convert to uint8 explicitly to avoid annoying errors
                            image = (255 * (image / image.max())).astype(np.uint8)

                        if volume_string not in oct_volume_dict:
                            oct_volume_dict[volume_string] = {}

                        if "image_index" in most_recent_bscan_metadata_dict:
                            oct_volume_dict[volume_string][
                                most_recent_bscan_metadata_dict["image_index"]
                            ] = {
                                "image": image,
                                "laterality": self.laterality,
                                "metadata": most_recent_bscan_metadata_dict,
                            }
                        else:
                            print(
                                f"missing image index from {most_recent_bscan_metadata_dict}"
                            )

                        if volume_string in volume_array_dict.keys():
                            volume_array_dict[volume_string][slice_index] = image
                        else:
                            print(
                                "Failed to save image data for volume {}".format(
                                    volume_string
                                )
                            )

            oct_volumes = []
            for key, volume in volume_array_dict.items():
                if len(volume) > 1:
                    # only use OCT volumes with more than 1 slice, the 1 slice volumes
                    # appear to be fundus photos
                    oct_volume = OCTVolumeWithMetaData(volume=volume, patient_id=key)
                    oct_volumes.append(oct_volume)
        return oct_volumes, oct_volume_dict, patient_metadata_dict

    def read_fundus_image(self):
        """Reads fundus data.
        Returns:
            obj:FundusImageWithMetaData
        """
        with open(self.filepath, "rb") as f:
            metadata_dict = {}
            raw = f.read(36)
            header = self.header_structure.parse(raw)

            raw = f.read(52)
            main_directory = self.main_directory_structure.parse(raw)

            # traverse list of main directories in first pass
            directory_stack = []
            laterality = None
            current = main_directory.current
            while current != 0:
                directory_stack.append(current)
                f.seek(current)
                raw = f.read(52)
                directory_chunk = self.main_directory_structure.parse(raw)
                current = directory_chunk.prev

            # traverse in second pass and  get all subdirectories
            chunk_stack = []
            for position in directory_stack:
                f.seek(position)
                raw = f.read(52)
                directory_chunk = self.main_directory_structure.parse(raw)

                for ii in range(directory_chunk.num_entries):
                    raw = f.read(44)
                    chunk = self.sub_directory_structure.parse(raw)
                    if chunk.start > chunk.pos:
                        chunk_stack.append([chunk.start, chunk.size])

            # initalise dict to hold all the image volumes
            image_array_dict = {}
            laterality_dict = {}
            # traverse all chunks and extract slices
            for start, pos in chunk_stack:
                f.seek(start)
                raw = f.read(60)
                try:
                    chunk = self.chunk_structure.parse(raw)
                except Exception as e:
                    continue

                if chunk.type == 11:  # laterality data
                    raw = f.read(20)
                    try:
                        laterality_data = self.lat_structure.parse(raw)
                        if laterality_data.laterality == 82:
                            laterality = "R"
                        elif laterality_data.laterality == 76:
                            laterality = "L"
                    except Exception:
                        laterality = None

                if chunk.type == 0x2714:  # bscan metadata
                    metadata_dict = self.get_metadata(f)

                if chunk.type == 0x40000000:  # image data
                    raw = f.read(self.image_structure.sizeof())
                    image_data = self.image_structure.parse(raw)
                    if chunk.ind == 0:  # fundus data
                        raw_volume = np.fromstring(
                            f.read(image_data.height * image_data.width), dtype=np.uint8
                        )
                        image = np.array(raw_volume).reshape(
                            image_data.height, image_data.width
                        )
                        image_string = "{}_{}_{}".format(
                            chunk.patient_id, chunk.study_id, chunk.series_id
                        )
                        image_array_dict[image_string] = image
                        laterality_dict[image_string] = laterality

            fundus_images = []
            for key, image in image_array_dict.items():
                fundus_images.append(
                    FundusImageWithMetaData(
                        image=image,
                        patient_id=key,
                        laterality=laterality_dict[key]
                        if key in laterality_dict.keys()
                        else None,
                    )
                )

        return fundus_images, metadata_dict
