# Creates fake photon event with a given pulse profile defined from time zero to one period
# The input is a FITS file with the event list as hdul[1]
# The output is a FITS file with the event list replaced and the profile data inserted into hdul[3].
from astropy.io import fits
import numpy as np
from scipy.interpolate import UnivariateSpline, RectBivariateSpline

def generate_event_file(input_spectrum,output_file,pulse_profile,ephemerides=None):
    with fits.open(input_spectrum) as hdul:
        chan=hdul[1].data['PI']
        eventtime=hdul[1].data['TIME']
        rawx=hdul[1].data['RAWX']
        rawy=hdul[1].data['RAWY']
        if "fits" in pulse_profile:
            with fits.open(pulse_profile) as hdu_profile:
                profile_hdu=hdu_profile['PROFILE']
                try:
                    nch=profile_hdu.header['NCHANNEL']
                    twod=True
                except KeyError:
                    twod=False
                if twod:
                    ntime=profile_hdu.header['NTIME']
                    ch=profile_hdu.data['CHANNEL'][0:nch]
                    tt=profile_hdu.data['TIME'][0:ntime]
                    rr=profile_hdu.data['RATE'].transpose()
                    rr=rr[0:nch,0:ntime]
                else:
                    tt=profile_hdu.data['TIME']
                    rr=profile_hdu.data['RATE']
        else:
            try:
                tt,rr=np.loadtxt(pulse_profile,usecols=(0,1),unpack=True)
                twod=False
            except IndexError:
                twod=True

            if twod:
                data=np.loadtxt(pulse_profile,usecols=(0),unpack=True)
                nch=int(data[0])
                ntime=int(data[1])
                size=int(nch*ntime)
                ch=data[2:nch+2]
                tt=data[nch+2:nch+ntime+2]
                rr=np.reshape(data[-size:],(nch,ntime))

        if twod:
            crr=np.cumsum(0.5*(rr[:,1:]+rr[:,:-1])*np.diff(tt),axis=1)
            crr=np.insert(crr,0,0,axis=1)
            crr=crr/np.tile(np.transpose([crr[:,-1]]),(1,ntime))
            clist=np.unique(crr)
            timearray=np.zeros((nch,len(clist)))
            #print(crr)
            #print(clist)
            for ii in range(nch):
                timearray[ii]=np.interp(clist,crr[ii],tt)
            #print(timearray)
            kx=nch-1
            if (kx>3):
                kx=3
            ky=ntime-1
            if (ky>3):
                ky=3;
            cfunk=RectBivariateSpline(ch,clist,timearray,s=0,kx=kx,ky=ky)
        else:
            ii=np.argsort(tt)
            tt=tt[ii]
            rr=rr[ii]
            crr=np.cumsum(0.5*(rr[1:]+rr[:-1])*np.diff(tt))
            crr=np.insert(crr,0,0)/crr[-1]
            cfunk=UnivariateSpline(crr,tt,s=0)

        rand=np.random.uniform(0,1,len(eventtime))

        if (ephemerides==None):
            if twod:
                for i in range(len(chan)):
                    eventtime[i]=np.floor(eventtime[i]/tt[-1])*tt[-1]+cfunk(chan[i],rand[i])
            else:
                eventtime=np.floor(eventtime/tt[-1])*tt[-1]+cfunk(rand)
        else:
            if "fits" in ephemerides:
                with fits.open(ephemerides) as hdu_ephem:
                    pulse_time=hdu_ephem['EPHEMERIDES'].data['TIME']
                    pulse_number=hdu_ephem['EPHEMERIDES'].data['PULSENUM']
            else:
                pulse_time,pulse_number=np.loadtxt(ephemerides,usecols=(0,1),unpack=True)
            pulsefunk=UnivariateSpline(pulse_time,pulse_number,s=0)
            timefunk=UnivariateSpline(pulse_number,pulse_time,s=0)
#            freqfunk=pulsefunk.derivative()
#            period=1/freqfunk(eventtime)
#            print(period)
#            eventtime=eventtime+(cfunk(rand)-np.modf(pulsefunk(eventtime))[1])*period
            if twod:
                for i in range(len(chan)):
                    eventtime[i]=timefunk(np.floor(pulsefunk(eventtime[i]))+cfunk(chan[i],rand[i])/tt[-1])
            else:
                eventtime=timefunk(np.floor(pulsefunk(eventtime))+cfunk(rand)/tt[-1])
            hdul.insert(3,fits.BinTableHDU.from_columns([
                fits.Column(name='TIME',format='E',unit='s',array=pulse_time),
                fits.Column(name='PULSENUM',format='E',array=pulse_number)
                ],name='EPHEMERIDES'))
            
        ii=np.argsort(eventtime)        
        hdul[1].data['PI']=chan[ii]
        hdul[1].data['TIME']=eventtime[ii]
        hdul[1].data['RAWX']=rawx[ii]
        hdul[1].data['RAWY']=rawy[ii]
        if twod:
            hdul.insert(3,fits.BinTableHDU.from_columns([
                fits.Column(name='TIME',format='E',unit='s',array=tt),
                fits.Column(name='CHANNEL',format='E',unit='s',array=ch),
                fits.Column(name='RATE',format='%dE' % nch,array=np.transpose(rr))
            ],name='PROFILE'))
            hdul[3].header['NCHANNEL']=(nch,'Number of energy channels in pulse profile')
            hdul[3].header['NTIME']=(ntime,'Number of times in pulse profile')
            hdul[3].header['COMMENT']='RATE array contains the folded spectrum at the given time'
        else:
            hdul.insert(3,fits.BinTableHDU.from_columns([
                fits.Column(name='TIME',format='E',unit='s',array=tt),
                fits.Column(name='RATE',format='E',array=rr)
            ],name='PROFILE'))
        hdul.writeto(output_file,output_verify='warn')
        # np.savetxt(output_file,np.transpose([t,channellist]),fmt=('%12.9f %d'))

    return

def main():
    import sys
    if len(sys.argv)<4:
        print("""Format:
python fakperiodic2.py _Input_FITS_Event_File_ _Output_FITS_Event_File_ _Pulse_Profile_ [_Ephemerides_]
""") 
    else:
        if len(sys.argv)>4:
            generate_event_file(sys.argv[1],sys.argv[2],sys.argv[3],sys.argv[4])
        else:
            generate_event_file(sys.argv[1],sys.argv[2],sys.argv[3])
            
    
# nbins=1000.0
# spec,b=np.histogram(plist, bins=nbins)
# plt.plot(b[1:],spec,'k-')
# plt.plot(chan,cnt*mf*nchan/nbins,'b-')
# plt.show()

if __name__ == "__main__":
    # execute only if run as a script
    main()