import os
import tempfile
import glob
import fitsio
import scoop
from scoop import utils

import settings


def extract_synth_table(filein, fileout):
    """
    Utility function to extract from spSpec?????-????-???_results_synth.fits file
    EMIS_SYNTH_PROF_DECOMP extension excluding arrays (LOSVD_*) and write into
    file out file.
    @param filein:
    @param fileout:
    """    
    # read file
    t = fitsio.read(filein, ext='EMIS_SYNTH_SPECTRA_DECOMP')
    # exclude columns with two dimentional arrays
    columns = ([n for n in t.dtype.names 
        if not n.startswith('LOSVD_') and 
        not n.startswith('SPEC_') and 
        not n.startswith('PARS_COMP_FLUX_') and
        not n.startswith('PARS_VSYS')])
    # write new table
    fitsio.write(fileout, t[columns], extname='restable_synth', clobber=True)


def get_spectrum_abspath(mjd, plate, fiberid, filename=None, search_dir=os.path.join(settings.SDSS_SPECTRA_PATH, 'fit_MILES_x')):
    """
    Utility function to look up absolute filename of a spectrum given its mjd, plate and fiberid.
    
    @param filename: if given, ignore mjd, plate, fiberid arguments and use this basename for lookup
    @return None: if not found
    
    Usage example:
    
    >>> get_spectrum_abspath(53818, 2013, 9, search_dir='/db2/Data/pro/SDSS/fit_MILES_x')
    '/db2/Data/pro/SDSS/fit_MILES_x/2013/spSpec53818-2013-009_results.fits'
    >>> get_spectrum_abspath(None, None, None, filename='spSpec53818-2013-009_results.fits', search_dir='/db2/Data/pro/SDSS/fit_MILES_x')
    '/db2/Data/pro/SDSS/fit_MILES_x/2013/spSpec53818-2013-009_results.fits'
    >>> get_spectrum_abspath(None, None, None, filename='/db2/Data/pro/SDSS/fit_MILES_x/2013/spSpec53818-2013-009_results.fits')
    '/db2/Data/pro/SDSS/fit_MILES_x/2013/spSpec53818-2013-009_results.fits'
    """
        
    if filename:
        # if input path exists, return its absolute path
        if os.path.exists(filename):
            return os.path.abspath(filename)
        # resolve non-absolute filename
        basename = os.path.basename(filename)
        plate = basename.split('-')[1]
        list = glob.glob(os.path.join(search_dir, "{}/{}".format(plate, basename)))
    else:
        # resolve by mjt, plate, fiberid
        list = glob.glob(os.path.join(search_dir, "{:04d}/spSpec{:05d}-{:04d}-{:03d}_results.fits".format(int(plate), int(mjd), int(plate), int(fiberid))))
        
    if len(list) == 1:
        return list[0]
    else:
        return None
    

def make_file_list(output_file=None, mask='fit_agn/04??', return_array=False):
    """
    Utility function to generate on the fly (temporary) file lists (well, SDSS spectra lists)
    
    @param output_file to create; if None, temporary file is created
    @param mask: which data subset to traverse, e.g. fit_agn/???? or fit_MILES_x/0101    
    @return: created filename with list of files inside
    """
    list = glob.glob(os.path.join(settings.SDSS_SPECTRA_PATH, mask, '*.fits'))
    if return_array:
        return list
    try:
        if output_file == None:
            output_file = tempfile.NamedTemporaryFile(delete=False)
            name = output_file.name
        else:
            name = output_file
            output_file = open(output_file, 'w')
        output_file.write("\n".join(list))
    finally:
        output_file.close()    
    
    return name


def make_tmp_file(prefix=None):
    """
    Make temporary file
    """
    if prefix == None:
        prefix = 'tmp'
    tmp_list = tempfile.NamedTemporaryFile(delete=False, prefix=prefix)
    tmp_list.close()
    return tmp_list.name


def write_to_tmp(list):
    """
    Wrapper on write to the temporary file function.
    @param: list: to be writen to the file
    @return: filename of temporary file. 
    """
    tmp_list = tempfile.NamedTemporaryFile(delete=False)
    tmp_list.write(list)
    tmp_list.close()
    return tmp_list.name


def log_head(args):
    """
    Make some information in the beginning of worker job.
    """
    scoop.logger.info("Data assemblying, worker {} at {} ({})".format(scoop.worker, utils.socket.gethostname(), utils.ip))
    scoop.logger.info("Current dir: {}".format(os.getcwd()))
    scoop.logger.info("ID: {}".format(args[0]+1) )
    scoop.logger.info("nIDs: {}".format(args[1]) )
    scoop.logger.info("Worker args: {}".format(args))
    scoop.logger.info("======================================================")
    return


if __name__ == "__main__":
    import doctest
    doctest.testmod()