import numpy as np
from astropy.io import fits
import msgpack
import msgpack_numpy as m
from poke.poke_core import Rayfront

m.patch()

def serialize(T):
    """serializes a class using msgpack
    written by Brandon Dube, docstring by Jaren Ashcraft

    Parameters
    ----------
    T : class
        class to convert to hex code. Used for rayfronts

    Returns
    -------
    serdat
        serial data corresponding to class T
    """
    glb = globals()
    Tname = T.__class__.__name__
    # assert Tname in glb, 'class must exist in globals in order to be re-hydrateable, with the same constraint'
    
    # now we make our storage format.  It will be:
    # 1) a header with the class name
    # 2) the content of vars(T)
    serdat = (Tname, vars(T))
    return msgpack.packb(serdat)

class MsgpackTrickerEmpty:
    """dummy class to trick msgpack
    """
    pass

def deserialize(buf):
    """deserializes a class using msgpack
    written by Brandon Dube, docstring by Jaren Ashcraft

    Parameters
    ----------
    buf : serdat
        serial data coorresponding to class

    Returns
    -------
    class
        deserialized class, typically a rayfront
    """
    e = MsgpackTrickerEmpty()
    Tname, varzzz = msgpack.unpackb(buf, use_list=True)
    for k, v in varzzz.items():
        setattr(e, k, v)
    
    e.__class__ = globals()[Tname]
    return e

def write_rayfront_to_serial(rayfront,filename):
    """writes rayfront to serial file using msgpack

    Parameters
    ----------
    rayfront : poke.Rayfront
        Rayfront object to serialize
    filename : str
        name of the file to save serial data to. The .msgpack extension will be added to this string
    """

    serdata = serialize(rayfront)

    with open(filename+'.msgpack','wb') as outfile:
        outfile.write(serdata)

def read_serial_to_rayfront(filename):
    """reads serial data containing Rayfront into a Rayfront object

    Parameters
    ----------
    filename : str
        name of the file to read serial data from

    Returns
    -------
    poke.Rayfront
        the saved poke.Rayfront object
    """

    with open(filename,'rb') as infile:
        serdata = infile.read()
    
    rayfront = deserialize(serdata)

    return rayfront


def jones_to_fits(data,filename,realimag=True):
    """Write a N x N x 2 x 2 x 2 Jones Pupil to a FITS file

    Parameters
    ----------
    data : numpy.ndarray complex128
        Jones pupil data generated by poke. Generally the first two dimensions are square, and correspond to the
        exit pupil coordinate. The remainder are matrix row, column, and then real/imag
    filename : str
        name of file to save the jones pupil as
    realimag : bool
        whether to save in real/imaginary parts. Defaults to True. If False, saves as absolute value and phase.
    """

    if realimag:
        realpart = np.real(data)
        imagpart = np.imag(data)
    else:
        realpart = np.abs(data)
        imagpart = np.angle(data)

    box = np.empty([data.shape[0],data.shape[1],2,2,2])
    box[...,0] = realpart
    box[...,1] = imagpart

    hdu = fits.PrimaryHDU(box)
    hdul = fits.HDUList([hdu])

    hdul['wavelength'] = 1.6
    hdul['pixelscale'] = 1

    pass