import numpy as np
import os
import subprocess as subp
import glob
import time
import re
import aztools as az
from ftplib import FTP
from astropy.io import fits as pyfits
import astropy.time as atime

base_dir = '/home/abzoghbi/data/ngc4151/spec_analysis'
data_dir = 'data/xmm'
os.system('mkdir -p %s'%data_dir)
obsids = ['0112310101', '0112310501', '0112830201', '0112830501', '0112830601',
          '0143500101', '0143500201', '0143500301', '0402660101', '0402660201',
          '0402660301', '0657840101', '0657840201', '0657840301', '0657840401',
          '0657840501', '0679780101', '0679780201', '0679780301', '0679780401',
          '0679780501', '0761670101', '0761670201', '0761670301', '0761670401',
          '0761670501', '0761670601', '0761670701', '0761670801', '0761670901']
obsids = np.array(obsids)
print('There are %d observations'%len(obsids))
print(', '.join(obsids))

os.chdir(base_dir)
ftp = FTP('legacy.gsfc.nasa.gov', 'anonymous', 'anonymous@gmail.com')
ftp.cwd('xmm/data/rev0')
failed = []
for o in obsids:
    tar_file = '%s/%s.tar'%(data_dir, o)
    # download file only if not already downloaded
    if not os.path.exists(tar_file):
        try:
            ftp.retrbinary('RETR %s.tar'%o ,open(tar_file, 'wb').write)
        except:
            print('failed downloading %s'%o)
            os.system('rm %s >/dev/null 2>&1'%tar_file)
            failed.append(o)

for f in failed:
    obsids = np.delete(obsids, np.argwhere(obsids==f)[0,0])
print('There are %d observations'%len(obsids))
print(', '.join(obsids))

os.chdir('%s/%s'%(base_dir, data_dir))
os.system('mkdir -p log')
procs = []
for o in obsids:
    if os.path.exists(o): continue
    os.system('tar -xf %s.tar'%o)
    os.chdir(o)
    os.system('rm -r 3XMM om_mosaic PPS >/dev/null 2>&1')
    os.system('mv ODF odf')
    os.chdir('odf')
    if not os.path.exists('ccf.cif'):
        os.system('gzip -d *gz')
        log_file = '../../log/%s_process.log'%o
        proc = subp.Popen(['/bin/bash', '-i', '-c', 'sasinit; xmm_process > %s 2>&1'%log_file])
        procs.append(proc)
    os.chdir('../..')

# wait for the tasks to end
for p in procs: p.wait()

os.chdir('%s/%s'%(base_dir, data_dir))
procs = []
for o in obsids:
    os.chdir(o)
    os.system('mkdir -p pn')
    os.chdir('pn')
    if len(glob.glob('*EVL*')) == 0 and len(glob.glob('pn.fits')) == 0:
        log_file = '../../log/%s_process_pn.log'%o
        p = subp.Popen(['/bin/bash', '-i', '-c', 'sasinit; xmm_process pn > %s 2>&1'%log_file])
        procs.append(p)
    os.chdir('../..')

# wait for the tasks to end
for p in procs: p.wait()

os.chdir('%s/%s'%(base_dir, data_dir))
exists = os.path.exists
for o in obsids:
    print('-- obs %s --'%o)
    os.chdir('%s/pn'%o)
    os.system('ln -s %s pn.fits'%glob.glob('*EVL*')[0])
    os.system('mkdir -p spec')
    os.chdir('spec')
    if not exists('pn_filtered.fits') or not exists('ds9.reg'):
        # check if we have a saved region file, or a temporary region file
        # for faster loading
        saved_reg = '../../../log/%s_ds9.reg'%o
        if exists(saved_reg):
            os.system('cp %s ds9.reg'%saved_reg)
            region = ''
        else:
            region = '--region'
            tmp_reg = '../../../log/ds9.reg'
            if exists(tmp_reg):
                os.system('cp %s tmp.reg'%tmp_reg)

        subp.call(['/bin/bash', '-i', '-c', 
                'sasinit; xmm_filter.py ../pn.fits pn --std %s'%region])
        if not exists(saved_reg):
            os.system('cp ds9.reg %s'%tmp_reg) 
    os.chdir('../../..')
    

no_data = ['0112310501', '0112830601', '0402660301', '0679780501', '0657840101']
obsids = np.array([o for o in obsids if not o in no_data])
print('There are %d observations'%len(obsids))
print(', '.join(obsids))

os.chdir('%s/%s'%(base_dir, data_dir))
print('obs_num | obsid | fractions | radius | pileup?')
for iobs, o in enumerate(obsids):
    if os.path.exists('%s/pn/pileup'%o): continue
    os.chdir('%s/pn'%o)
    os.system('mkdir -p pileup')
    os.chdir('pileup')
    
    pileup = True
    radius = 0.0
    reg_lines = open('../spec/ds9.reg').readlines()
    if 'background' in reg_lines[-1]:
        reg_lines = reg_lines[:-2] + [reg_lines[-1], reg_lines[-2]]
    g = re.match('\(.*\)\n', reg_lines[-1])
    dum_rad = reg_lines[-1].split(',')[2]
    reg_lines[-1] = reg_lines[-1].replace(',' + dum_rad + ',', ',%g,')
    reg_text = ''.join(reg_lines)
    

    while pileup:
        with open('ds9.reg', 'w') as fp: fp.write(reg_text%radius)
        subp.call(['/bin/bash', '-i', '-c', 
           'sasinit; xmm_spec.py ../pn.fits ds9.reg --check_pileup > pileup.log 2>&1'])
        line = [l for l in open('pileup.log').readlines() if '+/-' in l][0]
        frac = np.array(np.array(line.split())[[2,4,6,8]], np.double)
        pileup = (frac[0] > 1+3*frac[1]) or (frac[2] < 1-3*frac[3])
        # 1 arcsec = 20 pixels
        text = '{:4d} | {} | {} | {} | {:d}'.format(iobs+1, o, line[10:-1], radius/20., pileup)
        print(text)
        if pileup: radius = np.max([3.0*20, radius+10])
    os.chdir('../spec')
    os.system('cp ds9.reg ds9.reg.0')
    os.system('cp ../pileup/ds9.reg ds9.reg')
    os.chdir('../../../')
            

os.chdir('%s/%s'%(base_dir, data_dir))
procs = []
for iobs, o in enumerate(obsids):
    os.chdir('%s/pn/spec'%o)
    if len(glob.glob('spec*grp')) == 0:
        time.sleep(1)
        p = subp.Popen(['/bin/bash', '-i', '-c', 
           'sasinit; xmm_spec.py pn_filtered.fits ds9.reg -o spec_%d > spec.log 2>&1'%(iobs+1)])
        procs.append(p)
    os.chdir('../../..')
# wait for the tasks to end
for p in procs: p.wait()

os.chdir('%s/%s'%(base_dir, data_dir))
procs = []
for iobs, o in enumerate(obsids):
    os.chdir('%s/pn'%o)
    os.system('mkdir -p spec0')
    os.chdir('spec0')
    
    if len(glob.glob('spec0*grp')) == 0:
        # use region from spec, but with no central region extracion #
        reg_lines = open('../spec/ds9.reg').readlines()
        if 'background' in reg_lines[-1]:
            reg_lines = reg_lines[:-2] + [reg_lines[-1], reg_lines[-2]]
        g = re.match('\\\\(.*\\\\)\\\\n', reg_lines[-1])
        dum_rad = reg_lines[-1].split(',')[2]
        reg_lines[-1] = reg_lines[-1].replace(',' + dum_rad + ',', ',0,')
        with open('ds9.reg', 'w') as fp: fp.write(''.join(reg_lines))
        
    
        time.sleep(1)
        p = subp.Popen(['/bin/bash', '-i', '-c', 
           'sasinit; xmm_spec.py ../spec/pn_filtered.fits ds9.reg -o spec0_%d > spec.log 2>&1'%(iobs+1)])
        procs.append(p)
    os.chdir('../../..')
# wait for the tasks to end
for p in procs: p.wait()

os.chdir('%s/%s'%(base_dir, data_dir))
procs = []
for o in obsids:
    os.chdir(o)
    os.system('mkdir -p mos')
    os.chdir('mos')
    if len(glob.glob('*EVL*')) == 0:
        log_file = '../../log/%s_process_mos.log'%o
        time.sleep(1)
        p = subp.Popen(['/bin/bash', '-i', '-c', 'sasinit; xmm_process mos > %s 2>&1'%log_file])
        procs.append(p)
    os.chdir('../..')
# wait for the tasks to end
for p in procs: p.wait()

# wait until above is done! #
os.chdir('%s/%s'%(base_dir, data_dir))
exists = os.path.exists
for o in obsids:
    print('-- obs %s --'%o)
    os.chdir('%s/mos'%o)
    os.system('ln -s %s m1.fits'%glob.glob('*M1*EVL*')[0])
    os.system('ln -s %s m2.fits'%glob.glob('*M2*EVL*')[0])
    os.system('mkdir -p spec')
    os.chdir('spec')
    for m in ['m1', 'm2']:
        # this line comes after first run; these 2 obs have now m1 data
        if m == 'm1' and o in ['0657840201', '0761670401']: continue
        if not exists('%s_filtered.fits'%m) or not exists('ds9_%s.reg'%m):
            # check if we have a saved region file, or a temporary region file
            # for faster loading
            saved_reg = '../../../log/%s_ds9_%s.reg'%(o, m)
            if exists(saved_reg):
                os.system('cp %s ds9_%s.reg'%(saved_reg, m))
                region = ''
            else:
                region = '--region'
                tmp_reg = '../../pn/spec/ds9.reg'
                if exists(tmp_reg):
                    os.system('cp %s tmp.reg'%tmp_reg)

            subp.call(['/bin/bash', '-i', '-c', 
                    'sasinit; xmm_filter.py ../%s.fits %s --std %s'%(m, m, region)]) 
    os.chdir('../../..')

os.chdir('%s/%s'%(base_dir, data_dir))
print('obs_num | obsid | fractions | radius | pileup?')
for iobs, o in enumerate(obsids):
    os.chdir('%s/mos'%o)
    os.system('mkdir -p pileup')
    os.chdir('pileup')
    
    for m in ['m1', 'm2']:
        pileup = True
        radius = 0.0
        reg_file = '../spec/ds9_%s.reg'%m
        if not os.path.exists(reg_file): 
            print('no %s; skipping !'%reg_file)
            continue
        reg_lines = open(reg_file).readlines()
        if 'background' in reg_lines[-1]:
            reg_lines = reg_lines[:-2] + [reg_lines[-1], reg_lines[-2]]
        dum_rad = reg_lines[-1].split(',')[2]
        reg_lines[-1] = reg_lines[-1].replace(',' + dum_rad + ',', ',%g,')
        reg_text = ''.join(reg_lines)

        while pileup:
            with open('ds9_%s.reg'%m, 'w') as fp: fp.write(reg_text%radius)
            subp.call(['/bin/bash', '-i', '-c', 
               ('sasinit; xmm_spec.py ../%s.fits ds9_%s.reg --check_pileup'
                ' > pileup_%s.log 2>&1')%(m,m,m)])
            line = [l for l in open('pileup_%s.log'%m).readlines() if '+/-' in l][0]
            frac = np.array(np.array(line.split())[[2,4,6,8]], np.double)
            pileup = (frac[0] > 1+3*frac[1]) or (frac[2] < 1-3*frac[3])
            # 1 arcsec = 20 pixels
            text = '{:4d}-{} | {} | {} | {} | {:d}'.format(
                iobs+1, m, o, line[10:-1], radius/20., pileup)
            print(text)
            if pileup: radius = np.max([3.0*20, radius+10])
            if radius/20 > 6: break
    os.chdir('../spec')
    os.system('cp ds9_%s.reg ds9_%s.reg.0'%(m,m))
    os.system('cp ../pileup/ds9_%s.reg ds9_%s.reg'%(m,m))
    os.chdir('../../../')
            

os.chdir('%s/%s'%(base_dir, data_dir))
procs = []
print('mos 1')
for iobs, o in enumerate(obsids):
    os.chdir('%s/mos/spec'%o)
    if os.path.exists('ds9_m1.reg') and len(glob.glob('spec_m1*grp'))==0:
        p = subp.Popen(['/bin/bash', '-i', '-c', 
           ('sasinit; xmm_spec.py m1_filtered.fits ds9_m1.reg -o spec_m1_%d '
            '> spec_m1.log 2>&1')%(iobs+1)])
        procs.append(p)
        time.sleep(1)
    os.chdir('../../..')
# wait for the tasks to end
for p in procs: p.wait()

procs = []
print('mos 2')
for iobs, o in enumerate(obsids):
    os.chdir('%s/mos/spec'%o)
    if os.path.exists('ds9_m2.reg') and len(glob.glob('spec_m2*grp'))==0:
        p = subp.Popen(['/bin/bash', '-i', '-c', 
           ('sasinit; xmm_spec.py m2_filtered.fits ds9_m2.reg -o spec_m2_%d '
            '> spec_m2.log 2>&1')%(iobs+1)])
        procs.append(p)
        time.sleep(1)
    os.chdir('../../..')
# wait for the tasks to end
for p in procs: p.wait()

os.chdir('%s/%s'%(base_dir, data_dir))
procs = []
for o in obsids:
    os.chdir(o)
    os.system('mkdir -p rgs')
    os.chdir('rgs')
    if len(glob.glob('spec_rgs*')) == 0:
        log_file = '../../log/%s_process_rgs.log'%o
        p = subp.Popen(['/bin/bash', '-i', '-c', 'sasinit; xmm_process rgs > %s 2>&1'%log_file])
        time.sleep(1)
        procs.append(p)
    os.chdir('../..')
# wait for the tasks to end
for p in procs: p.wait()

os.chdir('%s/%s'%(base_dir, data_dir))
for iobs,o in enumerate(obsids):
    os.chdir('%s/rgs'%o)
    if len(glob.glob('spec_rgs_%d*'%iobs)) == 0:
        os.system('rename _rgs. _rgs_%d. spec_rgs*'%(iobs+1))
        root = 'spec_rgs_%d.'%(iobs+1) + '%s'
        with pyfits.open(root%'grp') as fp:
            fp[1].header['backfile'] = root%'bgd'
            fp[1].header['respfile'] = root%'rsp'
            os.system('rm tmp.grp > /dev/null 2>&1')
            fp.writeto('tmp.grp')
        os.system('mv %s _%s'%(root%'grp', root%'grp'))
        os.system('mv %s tmp.grp'%(root%'grp'))
        cmd = ('export HEADASNOQUERY=;export HEADASPROMPT=/dev/null;'
              'ftgrouppha tmp.grp %s opt respfile=%s')%(root%'grp', root%'rsp')
        subp.call(['/bin/bash', '-i', '-c', cmd])
    os.chdir('../..')

os.chdir('%s/%s/'%(base_dir, data_dir))
print('{:5} | {:12} | {:10.8} | {:10.3} | {:10.3}'.format(
        'num', 'obsid', 'mjd', 'rate', 'exposure'))
spec_data = []
for iobs,o in enumerate(obsids):
    with pyfits.open('%s/pn/spec/spec_%d.grp'%(o, iobs+1)) as fp:
        exposure = fp[1].header['exposure']
        counts = fp[1].data.field('counts').sum()
        mjd = (atime.Time(fp[0].header['date_end']).mjd + 
               atime.Time(fp[0].header['date_obs']).mjd ) / 2
        spec_data.append([mjd, counts/exposure, exposure/1e3])
        text = '{:5} | {:12} | {:10.8} | {:10.3} | {:10.3}'.format(
                iobs+1, o, mjd, counts/exposure, exposure/1e3)
        print(text)
spec_data = np.array(spec_data)

## keep only exposures > 5ks
igood = np.argwhere(spec_data[:,2] >= 5)[:,0]
spec_obsids = obsids[igood]
spec_data = spec_data[igood]
print('There are %d spec observations'%len(spec_obsids))
print(', '.join(spec_obsids))

os.chdir('%s/%s/'%(base_dir, data_dir))
# time_step in ks
time_step = 5
tselect, tselect_expr = [], []
print('{:5} | {:12} | {:10} | {:5}'.format('nu', 'obsid', 'exposure', 'nsub spec'))
for iobs,o in enumerate(spec_obsids):
    nsub_spec = np.int(spec_data[iobs, 2] // time_step)
    with pyfits.open('%s/pn/spec/pn_filtered.fits'%o) as fp:
        events = np.sort(fp[1].data.field('time'))
        evt_per_sub = len(events)//nsub_spec
        tcut = np.concatenate([[events[0]], events[np.arange(1, nsub_spec)*evt_per_sub], [events[-1]]])
        tselect_expr.append(['(TIME > %.10g) && (TIME <= %.10g)'%(x,y) 
                        for x,y in zip(tcut[:-1], tcut[1:])])
        tselect.append([[x,y] for x,y in zip(tcut[:-1], tcut[1:])])
    print('{:5} | {:12} | {:10.3} | {:5}'.format(iobs+1, o, spec_data[iobs, 2], nsub_spec))

os.chdir('%s/%s'%(base_dir, data_dir))
# use tselect array from above #
procs, ispec, ispecs = [], 1, []
for iobs in range(len(spec_obsids)):
    dum = [ispec+i for i in range(len(tselect[iobs]))]
    ispecs.append(dum)
    ispec = dum[-1]+1
        
for iobs,o in enumerate(spec_obsids):
    os.chdir('%s/pn'%o)
    os.system('mkdir -p subspec')
    os.chdir('subspec')

    if len(glob.glob('spec*grp')) != len(tselect_expr[iobs]):
        cmd = 'sasinit'
        for isel,tsel in enumerate(tselect_expr[iobs]):
            ispec = ispecs[iobs][isel]
            cmd += (';xmm_spec.py ../spec/pn_filtered.fits ../spec/ds9.reg'
                ' --e_expr " && %s" -o spec_%d > spec_%d.log 2>&1')%(tsel, ispec, ispec)
        time.sleep(0.5)
        p = subp.Popen(['/bin/bash', '-i', '-c', cmd])    
        procs.append(p)
    os.chdir('../../..')
# wait for the tasks to end
for p in procs: p.wait()

os.chdir('%s/%s'%(base_dir, data_dir))
# save some useful data for other notebooks
np.savez('log/data.npz', obsids=obsids, spec_obsids=spec_obsids, spec_data=spec_data, 
         tselect=tselect, tselect_ispec=ispecs)

base_dir = '/u/home/abzoghbi/data/ngc4151/spec_analysis'
data_dir = 'data/suzaku'
os.system('mkdir -p %s'%data_dir)
obsids = ['906006010', '906006020', '707024010', '701034010']
obsids = np.sort(obsids)
print('There are %d observations'%len(obsids))
print(', '.join(obsids))

os.chdir(base_dir)
ftp = FTP('legacy.gsfc.nasa.gov', 'anonymous', 'anonymous@gmail.com')
ftp.cwd('suzaku/data/obs')
failed = []
for o in obsids:
    tar_file = '%s/%s.tar'%(data_dir, o)
    ftp.cwd(o[0])
    # download file only if not already downloaded
    if not os.path.exists(tar_file):
        try:
            ftp.retrbinary('RETR %s.tar'%o ,open(tar_file, 'wb').write)
        except:
            print('failed downloading %s'%o)
            os.system('rm %s >/dev/null 2>&1'%tar_file)
            failed.append(o)
    ftp.cwd('..')

for f in failed:
    obsids = np.delete(obsids, np.argwhere(obsids==f)[0,0])
print('There are %d observations'%len(obsids))
print(', '.join(obsids))

os.chdir('%s/%s'%(base_dir, data_dir))
os.system('mkdir -p log')
procs = []
for o in obsids:
    if os.path.exists(o): continue
    os.system('tar -xf %s.tar'%o)
    if not os.path.exists('%s_p'%o):
        log_file = 'log/%s_process.log'%o
        cmd = ('export HEADASNOQUERY=;export HEADASPROMPT=/dev/null;'
               'suzaku_process %s xis > %s 2>&1'%(o, log_file))
        proc = subp.Popen(['/bin/bash', '-i', '-c', cmd])
        procs.append(proc)
        time.sleep(0.2)

# wait for the tasks to end
for p in procs: p.wait()

os.chdir('%s/%s'%(base_dir, data_dir))
exists = os.path.exists
procs = []
for iobs,o in enumerate(obsids):
    print('-- obs %s --'%o)
    os.chdir('%s_p/'%o)
    os.system('mkdir -p spec')
    os.chdir('spec')
    if len(glob.glob('spec_xi*grp')) != 3:
        # check if we have a saved region file, or a temporary region file
        # for faster loading
        saved_reg = '../../log/%s_src.reg'%o
        if exists(saved_reg):
            os.system('cp %s src.reg'%saved_reg)
            os.system('cp %s bgd.reg'%(saved_reg.replace('_src.', '_bgd.')))
            region = ''
        else:
            region = '--create_region'

        cmd = ('export HEADASNOQUERY=;export HEADASPROMPT=/dev/null;'
              'suzaku_xis_spec.py -o spec_%d %s'%(iobs+1, region))
        p = subp.Popen(['/bin/bash', '-i', '-c', cmd])
        procs.append(p)
        if not exists(saved_reg):
            os.system('cp src.reg %s'%saved_reg) 
            os.system('cp bgd.reg %s'%(saved_reg.replace('_src.', '_bgd.')))
        time.sleep(0.3)
    os.chdir('../..')
# wait for the tasks to end
for p in procs: p.wait() 

## add xi0, xi3 -> fi #
os.chdir('%s/%s'%(base_dir, data_dir))
for iobs,o in enumerate(obsids):
    os.chdir('%s_p/spec'%o)
    txt = '\n'.join(['spec_xi0_{0}.{1} spec_xi3_{0}.{1}'.format(iobs+1, s) 
                        for s in ['src', 'bgd', 'rsp']])
    with open('tmp.dat', 'w') as fp: fp.write(txt+'\n')
    cmd = ('export HEADASNOQUERY=;export HEADASPROMPT=/dev/null; rm spec_fi_{0}.*;'
            'addascaspec tmp.dat spec_fi_{0}.src spec_fi_{0}.rsp spec_fi_{0}.bgd').format(iobs+1)
    p = subp.call(['/bin/bash', '-i', '-c', cmd])
    # group #
    cmd = 'ogrppha.py spec_fi_{0}.src spec_fi_{0}.grp -f 3 -s 6'.format(iobs+1)
    subp.call(['/bin/bash', '-i', '-c', cmd])
    os.chdir('../..')

os.chdir('%s/%s/'%(base_dir, data_dir))
print('{:5} | {:12} | {:10.8} | {:10.8} | {:10.3} | {:10.3}'.format(
        'num', 'obsid', 'mjd_s', 'mjd_e', 'rate', 'exposure'))
spec_data = []
for iobs,o in enumerate(obsids):
    with pyfits.open('%s_p/spec/spec_xi0_%d.grp'%(o, iobs+1)) as fp:
        exposure = fp[1].header['exposure']
        counts = fp[1].data.field('counts').sum()
        tmid = np.array([fp[0].header['tstart'],  fp[0].header['tstop']])
        mref = fp[0].header['mjdrefi'] + fp[0].header['mjdreff']
        mjd = tmid / (24*3600) + mref
        spec_data.append([mjd, counts/exposure, exposure/1e3])
        text = '{:5} | {:12} | {:10.8} | {:10.8} | {:10.3} | {:10.5}'.format(
                iobs+1, o, mjd[0], mjd[1], counts/exposure, exposure/1e3)
        print(text)
spec_data = np.array(spec_data)

os.chdir('%s/%s/'%(base_dir, data_dir))
cmd_clean = 'export HEADASNOQUERY=;export HEADASPROMPT=/dev/null;'
# time_step in ks
time_step = 5
tselect, tselect_expr = [], []
print('{:5} | {:12} | {:10} | {:5}'.format('num', 'obsid', 'exposure', 'nsub spec'))
for iobs,o in enumerate(obsids):
    
    os.chdir('%s_p'%o)
    os.system('mkdir -p subspec')
    os.chdir('subspec')
    os.system('cp ../spec/*reg .')
    
    # extract light curve #
    lcfile = 'lc_008_xi0__1.src'
    if not os.path.exists(lcfile):
        cmd = cmd_clean + 'suzaku_xis_lc.py -t 8 --rootdir ../xis/event_cl'
        p = subp.call(['/bin/bash', '-i', '-c', cmd])

    ldata,dt = az.LCurve.read_fits_file(lcfile, min_exp=0.7)
    ldata = ldata[:, ldata[2]>0]
    nsub_spec = ldata.shape[1]//np.int(time_step * 1e3/dt)
    icut = np.array(np.linspace(0, ldata.shape[1]-1, nsub_spec+1), np.int)
    tcut = ldata[0, icut]
    with pyfits.open(lcfile) as fp:
        mjdref = fp[0].header['mjdrefi'] + fp[0].header['mjdreff']
    tcut = tcut/(24*3600) + mjdref
    
    tselect_expr.append(['mjd \\"%10.10g,%10.10g\\"'%(x,y) 
                        for x,y in zip(tcut[:-1], tcut[1:])])
    tselect.append([[x,y] for x,y in zip(tcut[:-1], tcut[1:])])
    print('{:5} | {:12} | {:10.5} | {:5}'.format(iobs+1, o, spec_data[iobs, 2], nsub_spec))
    os.chdir('../..')

os.chdir('%s/%s'%(base_dir, data_dir))
# use tselect array from above #
procs, ispec, ispecs = [], 1, []
for iobs in range(len(obsids)):
    dum = [ispec+i for i in range(len(tselect[iobs]))]
    ispecs.append(dum)
    ispec = dum[-1]+1
        
for iobs,o in enumerate(obsids):
    os.chdir('%s_p/subspec'%o)

    
    cmd0 = 'export HEADASNOQUERY=;export HEADASPROMPT=/dev/null;'
    for isel,tsel in enumerate(tselect_expr[iobs]):
        ispec = ispecs[iobs][isel]
        logfile = '../../log/sub_%d.log'%ispec
        cmd = cmd0 + 'suzaku_xis_spec.py -o spec_%d --t_expr "%s" --noclean > %s 2>&1'%(
                    ispec, tsel, logfile)

        if len(glob.glob('spec_xi*_%d.grp'%ispec)) != 3:
            time.sleep(1)
            p = subp.Popen(['/bin/bash', '-i', '-c', cmd])    
            procs.append(p)

        if len(procs) == 30:
            for p in procs: p.wait()
            procs = []
    
    os.chdir('../..')
# wait for the tasks to end
for p in procs: p.wait()

## add xi0, xi3 -> fi #
os.chdir('%s/%s'%(base_dir, data_dir))
for iobs,o in enumerate(obsids):
    os.chdir('%s_p/subspec'%o)
    cmd0 = 'export HEADASNOQUERY=;export HEADASPROMPT=/dev/null;'
    for isel,tsel in enumerate(tselect_expr[iobs]):
        ispec = ispecs[iobs][isel]
        if not os.path.exists('spec_fi_{0}.grp'.format(ispec)):
            txt = '\n'.join(['spec_xi0_{0}.{1} spec_xi3_{0}.{1}'.format(ispec, s) 
                            for s in ['src', 'bgd', 'rsp']])
            with open('tmp.dat', 'w') as fp: fp.write(txt)
            cmd = 'addascaspec tmp.dat spec_fi_{0}.src spec_fi_{0}.rsp spec_fi_{0}.bgd'.format(ispec)
            p = subp.call(['/bin/bash', '-i', '-c', cmd0+cmd])
            # group #
            cmd = 'ogrppha.py spec_fi_{0}.src spec_fi_{0}.grp -f 3 -s 6'.format(ispec)
            subp.call(['/bin/bash', '-i', '-c', cmd])
    os.chdir('../..')

os.chdir('%s/%s'%(base_dir, data_dir))
# save some useful data for other notebooks
np.savez('log/data.npz', obsids=obsids, spec_data=spec_data, 
         tselect=tselect, tselect_ispec=ispecs)

base_dir = '/u/home/abzoghbi/data/ngc4151/spec_analysis'
data_dir = 'data/nustar'
os.system('mkdir -p %s'%data_dir)
obsids = ['60001111002', '60001111003', '60001111005']
obsids = np.array(obsids)
print('There are %d observations'%len(obsids))
print(', '.join(obsids))

os.chdir('%s/%s'%(base_dir, data_dir))
ftp = FTP('legacy.gsfc.nasa.gov', 'anonymous', 'anonymous@gmail.com')
ftp.cwd('nustar/data/obs')
failed = []
for o in obsids:
    ftp.cwd('%s/%s'%(o[1:3], o[0]))
    tar_file = '%s.tar'%o
    # download file only if not already downloaded
    if not os.path.exists(tar_file):
        try:
            ftp.retrbinary('RETR %s'%tar_file ,open(tar_file, 'wb').write)
        except:
            print('failed downloading %s'%o)
            os.system('rm %s >/dev/null 2>&1'%tar_file)
            failed.append(o)
    ftp.cwd('../..')

for f in failed:
    if f in obsids:
        obsids = np.delete(obsids, np.argwhere(obsids==f)[0,0])
print('There are %d observations'%len(obsids))
print(', '.join(obsids))
nobs = len(obsids)

os.chdir('%s/%s'%(base_dir, data_dir))
os.system('mkdir -p log')
procs = []
for o in obsids:
    if not os.path.exists(o):
        os.system('tar -xf %s.tar'%o)
    if not os.path.exists('%s_p'%o):
        # download large files by http
        
        log_file = 'log/%s_process.log'%o
        cmd = ('export HEADASNOQUERY=;export HEADASPROMPT=/dev/null;'
               'nustar_process %s > %s 2>&1'%(o, log_file))
        proc = subp.Popen(['/bin/bash', '-i', '-c', cmd])
        procs.append(proc)
        time.sleep(0.2)

# wait for the tasks to end
for p in procs: p.wait()

os.chdir('%s/%s'%(base_dir, data_dir))
exists = os.path.exists
obsids = np.sort(obsids)
procs = []
for iobs,o in enumerate(obsids):
    print('-- obs %s --'%o)
    os.chdir('%s_p/'%o)
    os.system('mkdir -p spec')
    os.chdir('spec')
    if len(glob.glob('spec*grp')) != 2:
        # check if we have a saved region file, or a temporary region file
        # for faster loading
        saved_reg = '../../log/%s_src.reg'%o
        if exists(saved_reg):
            os.system('cp %s src.reg'%saved_reg)
            os.system('cp %s bgd.reg'%(saved_reg.replace('_src.', '_bgd.')))
            region = ''
        else:
            region = '--create_region'

        cmd = ('export HEADASNOQUERY=;export HEADASPROMPT=/dev/null;'
              'nustar_spec.py -o spec_%d %s'%(iobs+1, region))
        p = subp.Popen(['/bin/bash', '-i', '-c', cmd])
        procs.append(p)
        if not exists(saved_reg):
            os.system('cp src.reg %s'%saved_reg) 
            os.system('cp bgd.reg %s'%(saved_reg.replace('_src.', '_bgd.')))
        time.sleep(0.3)
    os.chdir('../..')
# wait for the tasks to end
for p in procs: p.wait() 

## group the spectra #
os.chdir('%s/%s'%(base_dir, data_dir))
for iobs,o in enumerate(obsids):
    os.chdir('%s_p/spec'%o)
    cmd = ('rm *grp; ogrppha.py spec_{0}_a_sr.pha spec_{0}_a.grp -f 3 -s 6;'
           'ogrppha.py spec_{0}_b_sr.pha spec_{0}_b.grp -f 3 -s 6').format(iobs+1)
    subp.call(['/bin/bash', '-i', '-c', cmd])
    os.chdir('../..')

os.chdir('%s/%s/'%(base_dir, data_dir))
print('{:5} | {:12} | {:10.8} | {:10.8} | {:10.3} | {:10.3}'.format(
        'num', 'obsid', 'mjd_s', 'mjd_e', 'rate', 'exposure'))
spec_data = []
for iobs,o in enumerate(obsids):
    with pyfits.open('%s_p/spec/spec_%d_a.grp'%(o, iobs+1)) as fp:
        exposure = fp[1].header['exposure']
        counts = fp[1].data.field('counts').sum()
        tmid = np.array([fp[0].header['tstart'], fp[0].header['tstop']])
        mref = fp[0].header['mjdrefi'] + fp[0].header['mjdreff']
        mjd = tmid / (24*3600) + mref
        spec_data.append([mjd, counts/exposure, exposure/1e3])
        text = '{:5} | {:12} | {:10.8} | {:10.8} | {:10.3} | {:10.5}'.format(
                iobs+1, o, mjd[0], mjd[1], counts/exposure, exposure/1e3)
        print(text)
spec_data = np.array(spec_data)

# summary #
os.chdir('%s/%s'%(base_dir, data_dir))
# save some useful data for other notebooks
np.savez('log/data.npz', obsids=obsids, spec_data=spec_data)

# ---- SUZAKU ---- #
# num   | obsid        | mjd_s      | mjd_e      | rat        | exp       
#     1 | 701034010    |   54087.84 |  54090.385 |       1.11 |     124.98
#     2 | 707024010    |  56242.862 |  56245.961 |       3.63 |     150.28
#     3 | 906006010    |  55882.674 |  55883.975 |        4.8 |     61.665
#     4 | 906006020    |   55913.67 |  55915.041 |       5.73 |     60.596

# ---- NUSTAR ---- #
# num   | obsid        | mjd_s      | mjd_e      | rat        | exp       
#     1 | 60001111002  |  56243.264 |  56243.765 |       7.64 |     21.864
#     2 | 60001111003  |  56243.792 |  56245.045 |       7.17 |     57.036
#     3 | 60001111005  |  56245.345 |   56246.73 |       8.26 |     61.531