import numpy as np
from construct import (Struct, Int16un, Int32un, PaddedString, Float64n, 
                        BytesInteger, Tell, Computed, Seek, Hex, Lazy, 
                        Array,this,Bytes)
from oct_converter.image_types import OCTVolumeWithMetaData
from pathlib import Path
import h5py, tempfile
                            
headerField = Struct(keylength=Int32un,
                     key=PaddedString(this.keylength, "utf8"), 
                     dataLength=Int32un)

floatField = Struct(keylength=Int32un,
                    key=PaddedString(this.keylength, "utf8"), 
                    dataLength=Int32un, 
                    value=Float64n)

intField = Struct(keylength=Int32un,
                  key=PaddedString(this.keylength, "utf8"), 
                  dataLength=Int32un, 
                  value=BytesInteger(this.dataLength,
                                     signed=False,
                                     swapped=True
                                     )
                  )

lazyIntField = Struct("keylength"/Int32un,
                        "key"/PaddedString(this.keylength, "utf8"), 
                        "dataLength"/Int32un, 
                        "offset"/Tell, 
                        "end"/ Computed(this.offset + this.dataLength),
                        "value" / Lazy(Bytes(this.dataLength)),
                        Seek(this.end)
            )

date = Struct(year=Int16un,month=Int16un,dow=Int16un,day=Int16un,hour=Int16un,minute=Int16un,second=Int16un,millisecond=Int16un)
dateField = Struct(keylength=Int32un,key=PaddedString(this.keylength, "utf8"), dataLength=Int32un, value=date)

strField = Struct(keylength=Int32un,key=PaddedString(this.keylength, "utf8"), dataLength=Int32un, value=PaddedString(this.dataLength, "utf8"))

class BOCT(object):
    """ Class for extracting data from Bioptigen's .OCT file format.
        .OCT stores 4d volumes (time series of 3D volumes with the same shape)

        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.
    """

    bioptigen_scan_type_map = {0:'linear',
                                1:'rect',
                                3:'rad'}

    def __init__(self, filepath):
        self.filepath = Path(filepath)
        if not self.filepath.exists():
            raise FileNotFoundError(self.filepath)

        self.oct_header = Struct(
                                "magicNumber" / Hex(Int32un),
                                "version" / Hex(Int16un),
                                "frameheader"/headerField,
                                "framecount"/ intField,
                                "linecount" / intField,
                                "linelength" / intField,
                                "sampleformat" / intField,
                                "description" / strField,
                                "xmin" / floatField,
                                "xmax" / floatField,
                                "xcaption" / strField,
                                "ymin" / floatField,
                                "ymax" / floatField,
                                "ycaption" / strField,
                                "scantype" / intField,
                                "scandepth" / floatField,
                                "scanlength" / floatField,
                                "azscanlength" / floatField,
                                "elscanlength" / floatField,
                                "objectdistance" / floatField,
                                "scanangle" / floatField,
                                "scans" / intField,
                                "frames" / intField,
                                "dopplerflag" / intField,
                                "config" / lazyIntField,
                                BytesInteger(4,signed=False,swapped=True)
                            )
        self.frame_header = Struct(
            "framedata" / headerField,
            "framedatetime" / dateField,
            "frametimestamp" / floatField,
            "framelines" / intField,
            "keylength"/Int32un,
            "key"/PaddedString(this.keylength, "utf8"), 
            "dataLength"/Int32un
            )

        self.frame_image = Struct(
            'rows' / Computed(this._._.header.linelength.value),
            'columns'/ Computed(this._.header.framelines.value),
            'totalpixels'/ Computed(this.rows*this.columns),
            'offset' / Tell,
            'end'/ Computed(this.offset + this.totalpixels*2),
            'pixels' / Lazy(Array(this.totalpixels,Int16un)),
            Seek(this.end)
            )

        self.frame = Struct(
            "header"/self.frame_header,
            "image" /self.frame_image,
            BytesInteger(4,signed=False,swapped=True)
            )
        
        self.frame_stack = Array(this.header.framecount.value, self.frame)
        self.file_structure = Struct(
                                "header" / self.oct_header,
                                "data" / self.frame_stack
                                )
    def read_oct_volume(self,diskbuffered=False):
        """ Reads OCT data.
            Args:
                diskbuffered (bool): If True, reduces memory usage by storing volume on disk using HDF5

            Returns:
                list(obj):[OCTVolumeWithMetaData]
        """
        ##Laterality/patient_id data not contained in .OCT file (often in filename)
        self.laterality = None
        self.patient_id = self.filepath.stem

        ##Lazily parse the file without loading frame pixels
        oct = self.file_structure.parse_file(self.filepath)
        header = oct.header
        self.frames = FrameGenerator(oct.data)
        scantype = self.bioptigen_scan_type_map[header.scantype.value]
        framecount = header.frames.value
        scancount = header.scans.value

        if scantype=='linear':
            ##linear bscans can contain multiple scans at one position
            ##reorder into (framecount+scancount,1,y,x)
            framecount += scancount
            scancount=1
        self.volume_shape = (framecount,scancount,self.frames.Ascans,self.frames.depth)
        self.vol_frames_shape = (self.volume_shape[0],self.volume_shape[1])
        
        ##Index conversion between 3D framestack (z+t,y,x) and 4D volume (t,z,x,y)
        stack_to_vol = np.asarray([i+j*(framecount) for i in range(0,framecount) for j in range(0,scancount)])
        vol_to_stack = np.asarray([i+j*(scancount) for i in range(0,scancount) for j in range(0,framecount)])
        self.indicies = {'to_vol':stack_to_vol,
                        'to_stack':vol_to_stack}
        self._loaded = False
        if diskbuffered:
            self._create_disk_buffer()
        else:
            self.vol = np.empty(self.volume_shape,dtype=np.uint16)

        return self.load_oct_volume()
    
    def _create_disk_buffer(self,name='vol'):
        _,_,x,y = self.volume_shape
        chunksize = (1,1,x,y)
        tf = h5py.File(tempfile.TemporaryFile(), "w")
        buffer = tf.create_dataset(name, 
                                    shape=self.volume_shape, 
                                    dtype = np.uint16,
                                    chunks=chunksize)
        setattr(self,name,buffer)

    def load_oct_volume(self):
        volFrames = np.reshape(self.frames.data,self.vol_frames_shape)
        try:
            with open(self.filepath, 'rb') as f:
                for t,v in enumerate(volFrames):
                    for z,frame in enumerate(v):
                        self.vol[t,z,:,:] = frame.load(f,self.frames.geom)
            self.loaded=True
        except Exception as e:
            print(e)
            print('Stopping load')
        return [OCTVolumeWithMetaData(self.vol[t,:,:,:]) for t in range(self.vol.shape[0])]

    def read_fundus_image(self):
        pass

class OCTFrame:
    def __init__(self,frame):
        self.frame = frame
        self.count = self.frame.image.totalpixels
        self.abs_pos = self.frame.image.offset

    def from_bytes(self,f):
        f.seek(self.abs_pos,0)
        im = np.fromfile(f, dtype=np.uint16, count=self.count)
        return im

    def load(self,f,imsize):
        return np.resize(self.from_bytes(f),imsize)

    def __eq__(self,other):
        return self.im == other


class FrameGenerator:
    def __init__(self,oct_data):
        self.data = None
        self.oct_data = oct_data

        frame0 = oct_data[0]
        self.Ascans = frame0.image.columns
        self.depth = frame0.image.rows
        self.geom = (self.Ascans, self.depth)
        self.data = np.asarray([OCTFrame(frame) for frame in self.oct_data])
        self.count = len(self.data)
        
    def set_frameorder(self,indexArr):
        try:
            self.data = self.data[indexArr]
        except Exception as e:
            print(e)