"""
"""

import os
import os.path as op


import datetime
import astropy.time
import pickle
try:
    import pyfits as pf
except ImportError:
    import astropy.io.fits as pf
    
OdometerFilename = "./odom"

class FitsBuilder:
    
    #EXPTYPES = ['BIAS', 'FLAT', 'DARK', 'LEDON', 'LEDOFF', 'STAR', 'IMG']
    EXPTYPES = ['SIM', 'IMG']
    
    def __init__(self, output_dir):
        """
        Builder. Maintains a list of (1) protoimages (i.e.  pixels +
        info read from the camera) and (2) meta-data collected from
        the XML-RPC servers, to be merged into a single header.
        """
        self.output_dir = output_dir
        if not op.isdir(output_dir):
            os.makedirs(output_dir)
        self.headers = {}
        self.pixels = []
        self.tables = {}
        
    def append(self, hdr_element, suffix='', prefix=''):
        """
        Append a header element.
        
        Parameters:
        -----------
          - hdr_element [dict of list]: hdr_element to append 
                        this is either a dict or a list of lists 
                        (cf Led XMLRPC server)
          - suffix, prefix [str]: suffix or prefix added to the dict keys
        """
        if(type(hdr_element) is type(None)):
            return
        elif type(hdr_element) is dict:
            keys, vals = list(hdr_element.keys()), list(hdr_element.values())
            keys = [prefix+k+suffix for k in keys]
            self.headers.update(dict(list(zip(keys, vals))))
        elif type(hdr_element) is list:
            for hdr_el in hdr_element:
                self.headers[prefix + hdr_el[0] + suffix] = str(hdr_el[1])
        else:
            raise ValueError('should be either dict or list')

    def add_image(self, data, exptype='IMG', **kw):
        """
        Append a new image to our internal proto_exposure list.
        """
        # make sure that imtype is valid
        if(type(data) is type(None)):
            return
        elif not exptype in FitsBuilder.EXPTYPES:
            raise ValueError('unknown exposure type: %r' % exptype)
        # generate a timestamp
        timestamp = datetime.datetime.utcnow()
        astropytime = astropy.time.Time(timestamp)
        # copy the data, because we are going to modify it
        my_data = data.copy()
        # dateobs and timestamps in the headers
        my_data['DATE-OBS'] = astropytime.fits
        my_data['MJD-OBS'] = astropytime.mjd
        # exposure type ? 
        my_data['EXPTYPE'] = exptype
        for k,v in list(kw.items()):
            my_data[k.upper()] = v
        # delete useless fields that may clash with pyfits
        try:
            del my_data['shape']
        except:
            pass
        self.pixels.append(my_data)

    def add_table(self, tablename, data):
        self.tables[tablename] = data
        
    def generate_filename(self, hdr, odo):
        """
        filename [str]: name of the file to be written
        """
        ret = ""
        
        # exposure type 
        ret += '%s_' % hdr['EXPTYPE']
        s = '%07d' % odo
        ret += s
             
        return ret + '.fits'
    
    
    def dump(self, output_dir=None):
        """
        Dump all the exposures in the internal exposure list.
        """
        if output_dir is None:
            output_dir = self.output_dir
        
        # general header, built from headera
        common_header = pf.Header(list(zip(list(self.headers.keys()), list(self.headers.values()))))
        with open(os.path.join(self.output_dir, 'headers.pkl'), 'ab') as fid:
            pickle.dump(self.headers, fid)
        curOdom = 0
        if op.isfile(OdometerFilename):
            curOdom = int(open(OdometerFilename, "r").read())
        filename=''
        if (len(self.pixels) == 0) and (len(self.tables) > 0):
            hh = pf.Header(common_header)
            hh.update({'EXPNUM': curOdom, 'EXPTYPE':'NOIMG'})
            hdul = pf.HDUList(pf.PrimaryHDU(header=hh))
            for tablename in self.tables:
                hdul.append(pf.TableHDU(data=self.tables[tablename], name=tablename))
            filename = output_dir + os.sep + self.generate_filename(hh, curOdom)
            hdul.writeto(filename)
            curOdom += 1
            
        for p in self.pixels:
            pix = p['pixels']
            del p['pixels']
            hh = pf.Header(list(zip(list(p.keys()), list(p.values()))))
            hh.update({'EXPNUM': curOdom})
            hh.update(common_header)
            hdul = pf.HDUList(pf.PrimaryHDU(header=hh, data=pix))
            if self.tables:
                for tablename in self.tables:
                    hdul.append(pf.TableHDU(data=self.tables[tablename], name=tablename))
            filename = output_dir + os.sep + self.generate_filename(hh, curOdom)
            hdul.writeto(filename)
            curOdom += 1
        
        outOdomFile = open(OdometerFilename, "w+")
        outOdomFile.write(str(curOdom))
        outOdomFile.close()
        return filename