# 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[2].
from astropy.io import fits
import numpy as np
from scipy.interpolate import UnivariateSpline

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']
        tt,rr=np.loadtxt(pulse_profile,usecols=(0,1),unpack=True)
        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)


        rand=np.random.uniform(0,crr[-1],len(eventtime))

        if (ephemerides==None):
            cfunk=UnivariateSpline(crr,tt,s=0)
            eventtime=np.floor(eventtime/tt[-1])*tt[-1]+cfunk(rand)
        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()
            tt=tt/tt[-1]
            cfunk=UnivariateSpline(crr,tt,s=0)
            period=1/freqfunk(eventtime)
#            print(period)
#            eventtime=eventtime+(cfunk(rand)-np.modf(pulsefunk(eventtime))[1])*period
            eventtime=timefunk(np.floor(pulsefunk(eventtime))+cfunk(rand))
            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]
        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 fakperiodic.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()