import numpy as np
import os
from astropy.io import fits, ascii
import PyROA
import signal

def handler(signum, frame):
        raise Exception("Timed out")

def runpyroa(rmid=np.arange(849),input_lag=None,Nsamples=15000,Nburnin=10000,fold=5.0,
             alllines=None,priors=None,init_tau=None,delay_dist=True,psi_types=None):

    # Set the timeout for the external routine
    timeout = 9000 # seconds


    # define initial setting
    topdir = '/data2/yshen/sdssrm/final_result/'
    output_dir = './pypetal/'

    if alllines is None:
        alllines =  ['ha', 'hb', 'mg2', 'c4']
    nline_all = len(alllines)
    doline = []
    for i in range(nline_all): 
        doline.append(False)

    nobj = len(rmid)
    if input_lag is None:
        file = '/data2/yshen/sdssrm/final_result/summary.fits'
        hdu = fits.open(file)
        tmp = hdu[1].data
        lagexp = tmp['LAG_OBS_EXP']
        hdu.close()
        input_lag = lagexp[rmid]

    for i in range(nobj):
        # setup path
        doline = [False,False,False,False]
        rmtag = "{:03d}".format(rmid[i])
        objdir = topdir + 'rm' + rmtag + '/'
        
        # check for existence of each line
        for j in range(nline_all):
            linefile = objdir + alllines[j] + '_lc.txt'
            if os.path.isfile(linefile):
                doline[j]=True

        useline = np.array(alllines)[np.where(doline)]
        nline = np.where(doline)[0].size

        if nline == 0:
            print('No line to fit: '+rmtag)
            continue

        # set up lag search range
        # this is the same search range for all nline lines
        lag0 = input_lag[i]
        if (lag0 < 1000):
            if (fold*lag0 < 2500.0):
                lag_bounds = [-fold*lag0, fold*lag0]
            else:
                lag_bounds = [-2500.0, 2500.0]
            if (fold*lag0 < 250.0):
                lag_bounds = [-250.0, 250.0]
        else:
            lag_bounds = [-2500.0, 4000.0]
      
        # Now do each line
        for iline in range(nline):
          
            laglimit = [lag_bounds]
            workdir = objdir + output_dir + useline[iline] + '/pyroa/'
            if os.path.exists(workdir) is False:
                os.makedirs(workdir)
            os.chdir(workdir)

            lcdir = objdir + output_dir + 'processed_lcs/'
            # setup pyroa input light curves (cont outlier rejected)
            # note that I normalize the light curves to have mean=1
            contfile = lcdir + 'pyroa_cont.dat'
            if os.path.isfile(contfile) is False:
                file0 = lcdir+ 'cont_data.dat'
                xi, yi, erri = np.loadtxt(file0, delimiter=',', unpack=True, usecols=[0,1,2])
                meanflux = np.mean(yi)
                yi, erri = yi/meanflux, erri/meanflux
                ascii.write( [xi, yi, erri] , contfile, format='no_header' )
                # np.savetxt(contfile, np.array([xi, yi, erri]) , fmt='%0.3f')
            linefile = lcdir + 'pyroa_' + useline[iline] + '.dat'
            if os.path.isfile(linefile) is False:
                file0 = lcdir + useline[iline] + '_data.dat'
                xi, yi, erri = np.loadtxt(file0, delimiter=',', unpack=True, usecols=[0,1,2])
                meanflux = np.mean(yi)
                yi, erri = yi/meanflux, erri/meanflux
                ascii.write( [xi, yi, erri], linefile, format='no_header')
                # np.savetxt(linefile, np.array([xi, yi, erri]) , fmt='%0.3f')

            # now run pyroa
            if priors is None:
                priors = [[0.0, 2.0],[0.0, 2.0], lag_bounds, [5.0, 50.0], [0.0, 10.0]]
            if init_tau is None:
                init_tau = [10]
            filters = ['cont', useline[iline]]

            # Call external routine
            try:
               signal.signal(signal.SIGALRM, handler)
               signal.alarm(timeout)
               # now run pyroa
               fit = PyROA.Fit(lcdir, 'pyroa', filters, priors, 
                               Nsamples = Nsamples, Nburnin = Nburnin,psi_types=psi_types, 
                               add_var=True, delay_dist=delay_dist, init_delta=10, init_tau=init_tau) # for the first pass I used init_tau=10
               signal.alarm(0)  # Cancel the timeout
            except Exception as exc:
               print(f"External routine timed out for {obj}: {exc}")
               with open(topdir+"killed_objects.txt", "a") as f:
                       f.write(str(obj) + "\n")
               continue

            print('Finished: '+rmtag + ' line: '+useline[iline])