#!/usr/bin/env python3
import argparse
import numpy
from astropy.io import fits as pyfits
import PyParadise

__author__ = "Bernd Husemann"
__contributor__ = ['Bernd Husemann', 'Omar Choudhury']
__credit__ = ['Bernd Husemann', 'Omar Choudhury','Jakbo Walcher','Tanya Urrutia','Anika Beer']
__copyright__ = "Copyright 2020, Bernd Husemann"
__license__ = "MIT"
__maintainer__ = "Bernd Husemann"
__email__ = "berndhusemann@gmx.de"
__status__ = "Development"
__version__ = "0.2"

class ParadiseApp(object):
    """The class `ParadiseApp` handles the fitting of the spectrum with
    a linear combination of template spectra and/or emission lines. For the
    the fitting of the template spectra, velocity and the velocity
    dispersion are obtained by Monte Carlo Markov Chains and the linear
    combination of template spectra with a non-negative least-squares
    algorithm.

    The emission line fitting is a separate function, and fits a provided
    set of emission lines.

    The bootstrapping fitting function can be used to obtain errors on the
    parameters for the template spectral fitting and errors on the
    parameters of the emission lines.
    """
    def __init__(self, input_file, outprefix, instrFWHM):
        """
        Parameters
        ----------
        input_file : str
            The fits file containing the spectrum for which the fitting will be
            performed
        outprefix : str
            The prefix that will be prepended to the files that are written by
            the fitting routines.
        instrFWHM : float
            The instrumental resolution of the data in `input_file`.
        """

        self.__inputData = PyParadise.spectrum1d.loadSpectrum(input_file)
        if self.__inputData._datatype == 'CUBE':
            self.__inputData = PyParadise.cube.loadCube(input_file)
            self.__inputData.correctError()
            self.__datatype = 'CUBE'
        elif self.__inputData._datatype == 'RSS':
            self.__inputData = PyParadise.rss.loadRSS(input_file)
            self.__inputData.correctError()
            self.__datatype = 'RSS'
        elif self.__inputData._datatype == 'Spectrum1D':
            data = PyParadise.spectrum1D.loadSpectrum(input_file)
            data.correctError()
            self.__datatype = 'RSS'
            if data._error is not None:
                err = numpy.array([data._error])
            else:
                err = None
            if data._mask is not None:
                m = numpy.array([data._mask])
            else:
                m = None
            self.__inputData = PyParadise.rss.RSS(wave=data._wave, data=numpy.array([data._data]), error=err, mask=m)
        self.__outPrefix = outprefix
        try:
            self.__instrFWHM = PyParadise.spectrum1d.SpectralResolution(res=float(instrFWHM))
        except ValueError:
            try:
                self.__instrFWHM = PyParadise.spectrum1d.SpectralResolution()
                self.__instrFWHM.readFile(instrFWHM)
            except IOError:
                print("Wrong input for spectral resolution. Specify either a float number or the path to an ASCII file with columns wavelengt and spectral resolution as content.")

    def run_SSP_fit(self, parfile, parallel, verbose):
        """This functions fits a linear combination of template spectra to the
        input spectra to obtain the best fit.

        Parameters
        ----------
        parfile : str
            The parameter file containing the constraints under which the input
            spectra will be fitted.
        parallel : {'auto', int}, optional
            If parallel is not equal to one, the python multiprocessing routine
            shall be used to run parts of the code in parallel. With the option
            `auto`, it adjusts the number of parallel processes to the number
            of cpu-cores/threads that are available.
        verbose : bool, optional
            Produces screen output, such that the user can follow which spectrum
            is currently fitted and what the results are for each spectrum.

        Notes
        -----
        The function does not return anything, but writes the results of the fit
        to disk. There are three files written to disk, all prepended by the
        `outprefix` supplied to this class at creation. The three files are:
        - `outprefix`.cont_model.fits which contains the spectra corresponding
          to the best linear combination of spectra.
        - `outprefix`.cont_res.fits contains the residual between the input fits
          file and the best fitted model spectra.
        - `outprefix`.stellar_table.fits contains the parameters of the best
          fit, like velocity, velocity dispersion, the luminosity-weighted and
          mass-weighted parameters.
        """
        ## Read in parameter file and get values
        parList = PyParadise.parameters.ParameterList(parfile)
        nlib_guess = int(parList['tmplinitspec'].getValue())
        vel_guess = float(parList['vel_guess'].getValue())
        vel_min = float(parList['vel_min'].getValue())
        vel_max = float(parList['vel_max'].getValue())
        disp_min = float(parList['disp_min'].getValue())
        disp_max = float(parList['disp_max'].getValue())
        kin_fix = bool(int(parList['kin_fix'].getValue()))
        mcmc_code = parList['mcmc_code'].getValue()
        nwidth_norm = int(parList['nwidth_norm'].getValue())
        iterations = int(parList['iterations'].getValue())
        samples = int(parList['samples'].getValue())
        walkers = int(parList['walkers'].getValue())
        burn = int(parList['burn'].getValue())
        thin = int(parList['thin'].getValue())
        store_chain = bool(int(parList['store_chain'].getValue()))
        start_wave = float(parList['start_wave'].getValue())
        end_wave = float(parList['end_wave'].getValue())
        excl_fit = PyParadise.parameters.CustomMasks(parList['excl_fit'].getValue())
        excl_cont = PyParadise.parameters.CustomMasks(parList['excl_cont'].getValue())
        min_x = int(parList['min_x'].getValue())
        max_x = int(parList['max_x'].getValue())
        min_y = int(parList['min_y'].getValue())
        max_y = int(parList['max_y'].getValue())
        bins = [[None, None]]

        try:
            f = open(parList['agebinfile'].getValue())
            try:
                for l in f.readlines():
                    bins.append([float(l.split()[0]), float(l.split()[1])])
            finally:
                f.close()
        except (IOError, ValueError):
            if verbose:
                print('agebinfile is non-existent or ignored.')
                print('Optional stellar population binning will not occur.')

        if verbose:
            print("The stellar population library is being prepared.")
        lib = PyParadise.ssplibrary.SSPlibrary(filename=parList['tmpldir'].getValue() + '/' + parList['tmplfile'].getValue())
        if kin_fix is True:
            hdu = pyfits.open(self.__outPrefix + '.kin_table.fits')
            tab = hdu[1].data
            vel_fit = tab.field('vel_fit')
            disp_fit = tab.field('disp_fit')
            vel_fit_err = tab.field('vel_fit_err')
            disp_fit_err = tab.field('disp_fit_err')
            Rvel = -numpy.ones(vel_fit.shape)
            Rdisp = -numpy.ones(disp_fit.shape)
            if self.__datatype == 'CUBE':
                x_pixels = tab.field('x_cor')
                y_pixels = tab.field('y_cor')
            elif self.__datatype == 'RSS':
                fibers = tab.field('fiber')
            vel_min = numpy.min(vel_fit)
            vel_max = numpy.max(vel_fit)

        min_wave = (start_wave / (1 + (vel_max + 2000) / 300000.0))
        max_wave = (end_wave / (1 + (vel_min - 2000) / 300000.0))
        if nlib_guess < 0:
            select = numpy.arange(lib.getBaseNumber()) == nlib_guess * -1 - 1
            lib = lib.subLibrary(select)
        lib = lib.subWaveLibrary(min_wave=min_wave, max_wave=max_wave)
        lib = lib.matchInstFWHM(self.__instrFWHM, vel_guess)
        lib = lib.resampleWaveStepLinear(self.__inputData.getWaveStep(), vel_guess / 300000.0)
        lib_norm = lib.normalizeBase(nwidth_norm, excl_cont, vel_guess / 300000.0)
        lib_rebin = lib_norm.rebinLogarithmic()

        if verbose:
            print("The input cube is being normalized.")

        normData = self.__inputData.normalizeSpec(nwidth_norm, excl_cont.maskPixelsObserved(self.__inputData.getWave(),
             vel_guess / 300000.0))
        normDataSub = normData.subWaveLimits(start_wave, end_wave)

        #pylab.plot(self.__inputData._data[0,:],'-k')
        #pylab.plot(normData._normalization[0,:],'-g')
        #pylab.plot(normData._data[0,:],'-r')
        #pylab.show()
        if verbose:
            print("The stellar population modelling has been started.")
        if self.__datatype == 'CUBE':
            if kin_fix:
                (fitted, coeff, chi2, x_pix, y_pix, cube_model, mask) = normDataSub.fit_Lib_fixed_kin(lib_rebin, nlib_guess,
                vel_fit, disp_fit, x_pixels, y_pixels, min_x, max_x, min_y, max_y,
                excl_fit, verbose, parallel)
            else:
                (vel_fit, vel_fit_err, Rvel, disp_fit, disp_fit_err, Rdisp, fitted, coeff, chi2, x_pix, y_pix,
                cube_model, mask) = normDataSub.fit_Kin_Lib_simple(lib_rebin, nlib_guess, vel_min, vel_max, disp_min, disp_max,
                min_x, max_x, min_y, max_y, excl_fit, iterations, mcmc_code,
                walkers, burn, samples, thin, verbose, parallel)
        elif self.__datatype == 'RSS':
            if kin_fix:
                (fitted, coeff, chi2, fiber, rss_model, mask) = normDataSub.fit_Lib_fixed_kin(lib_rebin, nlib_guess, vel_fit, disp_fit, fibers, min_y, max_y, excl_fit, verbose, parallel)
            else:

                (vel_fit, vel_fit_err, Rvel, disp_fit, disp_fit_err, Rdisp, fitted, coeff, chi2, fiber, rss_model, mask, vel_trace, disp_trace) = normDataSub.fit_Kin_Lib_simple(lib_rebin, nlib_guess,vel_min, vel_max, disp_min, disp_max, min_y, max_y, excl_fit, iterations, mcmc_code, walkers, burn, samples, thin, verbose, store_chain, parallel)
        if verbose:
                print("Storing the results to %s (model), %s (residual) and %s (parameters)." % (
                    self.__outPrefix + '.cont_model.fits', self.__outPrefix + '.cont_res.fits',
                    self.__outPrefix + '.stellar_table.fits'))
        # pylab.plot(rss_model[0,:],'-g')
        # pylab.show()
        if self.__datatype == 'RSS':
            model_out = PyParadise.rss.RSS(wave=normDataSub.getWave(), data=rss_model, error=normDataSub.unnormalizedSpec()._error, mask=mask,
                header=self.__inputData.getHeader(), normalization=normDataSub.getNormalization())
            res_out = PyParadise.rss.RSS(wave=normDataSub.getWave(),
                data=self.__inputData.subWaveLimits(start_wave, end_wave).getData() - rss_model,
                header=self.__inputData.getHeader())
        elif self.__datatype == 'CUBE':
            model_out = PyParadise.cube.Cube(wave=normDataSub.getWave(), data=cube_model, error=normDataSub.unnormalizedSpec()._error, mask=mask,
                header=self.__inputData.getHeader(), normalization=normDataSub.getNormalization())
            res_out = PyParadise.cube.Cube(wave=normDataSub.getWave(),
                data=self.__inputData.subWaveLimits(start_wave, end_wave).getData() - cube_model,
                header=self.__inputData.getHeader())
        if numpy.max(self.__inputData.getWave()[1:] - self.__inputData.getWave()[:-1]) - numpy.min(
            self.__inputData.getWave()[1:] - self.__inputData.getWave()[:-1]) < 0.01:
            model_out.writeFitsData(self.__outPrefix + '.cont_model.fits')
            res_out.writeFitsData(self.__outPrefix + '.cont_res.fits')
        else:
            model_out.writeFitsData(self.__outPrefix + '.cont_model.fits', store_wave=True)
            res_out.writeFitsData(self.__outPrefix + '.cont_res.fits', store_wave=True)

        mass_weighted_pars = numpy.zeros((len(fitted), 5, len(bins) + 1), dtype=numpy.float32)
        lum_weighted_pars = numpy.zeros((len(fitted), 5, len(bins) + 1), dtype=numpy.float32)
        for i in range(len(fitted)):
            if fitted[i]:
                for j in range(len(bins)):
                    try:
                        mass_weighted_pars[i, :, j] = lib_norm.massWeightedPars(coeff[i, :], bins[j][0], bins[j][1])
                    except:
                        mass_weighted_pars[i, :, j] = numpy.array([numpy.nan, numpy.nan, numpy.nan, numpy.nan, numpy.nan])
                    try:
                        lum_weighted_pars[i, :, j] = lib_norm.lumWeightedPars(coeff[i, :], bins[j][0], bins[j][1])
                    except:
                        lum_weighted_pars[i, :, j] = numpy.array([numpy.nan, numpy.nan, numpy.nan, numpy.nan, numpy.nan])

        columns = []
        if self.__datatype == 'CUBE':
            columns.append(pyfits.Column(name='x_cor', format='J', array=x_pix[fitted]))
            columns.append(pyfits.Column(name='y_cor', format='J', array=y_pix[fitted]))
        elif self.__datatype == 'RSS':
            columns.append(pyfits.Column(name='fiber', format='J', array=fiber[fitted]))
        columns.append(pyfits.Column(name='vel_fit', format='E', unit='km/s', array=vel_fit[fitted]))
        columns.append(pyfits.Column(name='vel_fit_err', format='E', unit='km/s', array=vel_fit_err[fitted]))
        if store_chain and self.__datatype == 'RSS':
            columns.append(pyfits.Column(name='vel_trace', format='%dE'%(vel_trace.shape[1]), unit='km/s', array=vel_trace[fitted,:]))
        columns.append(pyfits.Column(name='Rvel', format='E', unit='km/s', array=Rvel[fitted]))
        columns.append(pyfits.Column(name='disp_fit', format='E', unit='km/s', array=disp_fit[fitted]))
        columns.append(pyfits.Column(name='disp_fit_err', format='E', unit='km/s', array=disp_fit_err[fitted]))
        if store_chain and self.__datatype == 'RSS':
            columns.append(pyfits.Column(name='disp_trace', format='%dE'%(disp_trace.shape[1]), unit='km/s', array=disp_trace[fitted,:]))
        columns.append(pyfits.Column(name='Rdisp', format='E', unit='km/s', array=Rdisp[fitted]))
        columns.append(pyfits.Column(name='chi2', format='E', array=chi2[fitted]))
        if lib.getBaseNumber() > 1:
            columns.append(pyfits.Column(name='base_coeff', format='%dE' % (lib.getBaseNumber()), array=coeff[fitted, :]))
        else:
            columns.append(pyfits.Column(name='base_coeff', format='E', array=coeff[fitted].flatten()))
        for i, postfix in enumerate(['total'] + ['bin{}'.format(i) for i in range(1, len(bins))]):
            columns.append(pyfits.Column(name='lum_coeff_frac_' + postfix, format='E', array=lum_weighted_pars[fitted, 0, i]))
            columns.append(pyfits.Column(name='lum_age_' + postfix, format='E', array=lum_weighted_pars[fitted, 1, i]))
            columns.append(pyfits.Column(name='lum_M/L_' + postfix, format='E', array=lum_weighted_pars[fitted, 2, i]))
            columns.append(pyfits.Column(name='lum_[Fe/H]_' + postfix, format='E', array=lum_weighted_pars[fitted, 3, i]))
            columns.append(pyfits.Column(name='lum_[A/Fe]_' + postfix, format='E', array=lum_weighted_pars[fitted, 4, i]))
            columns.append(pyfits.Column(name='mass_coeff_frac_' + postfix, format='E', array=mass_weighted_pars[fitted, 0, i]))
            columns.append(pyfits.Column(name='mass_age_' + postfix, format='E', array=mass_weighted_pars[fitted, 1, i]))
            columns.append(pyfits.Column(name='mass_M/L_' + postfix, format='E', array=mass_weighted_pars[fitted, 2, i]))
            columns.append(pyfits.Column(name='mass_[Fe/H]_' + postfix, format='E', array=mass_weighted_pars[fitted, 3, i]))
            columns.append(pyfits.Column(name='mass_[A/Fe]_' + postfix, format='E', array=mass_weighted_pars[fitted, 4, i]))

        try:
            table_out = pyfits.BinTableHDU.from_columns(columns)
        except:
            table_out = pyfits.new_table(columns)
        table_out.writeto(self.__outPrefix + '.stellar_table.fits', overwrite=True)

    def run_eline_fit(self, parfile, parallel, verbose):
        """This functions fits a a set of emission lines to the spectra.

        Parameters
        ----------
        parfile : str
            The parameter file containing the constraints under which the
            emission lines will be fitted to the input spectra.
        parallel : {'auto', int}, optional
            If parallel is not equal to one, the python multiprocessing routine
            shall be used to run parts of the code in parallel. With the option
            `auto`, it adjusts the number of parallel processes to the number
            of cpu-cores/threads that are available.
        verbose : bool, optional
            Produces screen output, such that the user can follow which spectrum
            is currently fitted and what the results are for each spectrum.

        Notes
        -----
        The function does not return anything, but writes the results of the fit
        to disk. There are three files written to disk, all prepended by the
        `outprefix` supplied to this class at creation. The three files are:
        - `outprefix`.eline_model.fits which contains the best fitted emission
          line spectra.
        - `outprefix`.eline_res.fits contains the residual between the input
          fits file and the best fitted emission line model spectra.
        - `outprefix`.eline_table.fits contains the parameters of the best
          fit, like flux, velocity and the FWHM.
        """
        parList = PyParadise.parameters.ParameterList(parfile)
        eCompFile = parList['eCompFile'].getValue()
        vel_guess = float(parList['vel_guess'].getValue())
        line_fit = PyParadise.parameters.CustomMasks(parList['line_fit_region'].getValue())
        efit_method = parList['efit_method'].getValue()
        efit_ftol = float(parList['efit_ftol'].getValue())
        efit_xtol = float(parList['efit_xtol'].getValue())
        guess_window = int(parList['eguess_window'].getValue())
        min_x = float(parList['min_x'].getValue())
        max_x = float(parList['max_x'].getValue())
        min_y = float(parList['min_y'].getValue())
        max_y = float(parList['max_y'].getValue())

        hdu = pyfits.open(self.__outPrefix + '.stellar_table.fits')
        stellar_table = hdu[1].data
        if self.__datatype == 'CUBE':
            x_cor = stellar_table.field('x_cor')
            y_cor = stellar_table.field('y_cor')
        elif self.__datatype == 'RSS':
            fiber = stellar_table.field('fiber')

        line_par = PyParadise.fit_profile.parFile(eCompFile, self.__instrFWHM)
        if self.__datatype == 'CUBE':
            res_out = PyParadise.cube.loadCube(self.__outPrefix + '.cont_res.fits')
        elif self.__datatype == 'RSS':
            res_out = PyParadise.rss.loadRSS(self.__outPrefix + '.cont_res.fits')
        disp1 = res_out._wave[1]-res_out._wave[0]
        res_wave_start = res_out._wave[0]-disp1/2.0
        disp2 = res_out._wave[-1]-res_out._wave[-2]
        res_wave_end = res_out._wave[-1]+disp2/2.0
        res_out._error = self.__inputData.subWaveLimits(res_wave_start, res_wave_end)._error
        res_out._mask = self.__inputData.subWaveLimits(res_wave_start, res_wave_end)._mask
        if self.__datatype == 'CUBE':
            out_lines = res_out.fitELines(line_par, line_fit.maskPixelsObserved(res_out.getWave(),
                vel_guess / 300000.0), min_x, max_x, min_y, max_y, method=efit_method, guess_window=guess_window,
                spectral_res=self.__instrFWHM, ftol=efit_ftol, xtol=efit_xtol, verbose=verbose, parallel=parallel)
            model_line = PyParadise.cube.Cube(wave=res_out.getWave(), data=out_lines[4], header=self.__inputData.getHeader())
            line_res = PyParadise.cube.Cube(wave=res_out.getWave(), data=res_out.getData() - model_line.getData(),
                header=self.__inputData.getHeader())
        elif self.__datatype == 'RSS':
            out_lines = res_out.fitELines(line_par, line_fit.maskPixelsObserved(res_out.getWave(),
                vel_guess / 300000.0), min_y, max_y, method=efit_method, guess_window=guess_window,
                spectral_res=self.__instrFWHM, ftol=efit_ftol, xtol=efit_xtol, verbose=verbose, parallel=parallel)
            model_line = PyParadise.rss.RSS(wave=res_out.getWave(), data=out_lines[3], header=self.__inputData.getHeader())
            line_res = PyParadise.rss.RSS(wave=res_out.getWave(), data=res_out.getData() - model_line.getData(),
                header=self.__inputData.getHeader())
        if numpy.max(self.__inputData.getWave()[1:] - self.__inputData.getWave()[:-1]) - numpy.min(
                     self.__inputData.getWave()[1:] - self.__inputData.getWave()[:-1]) < 0.01:
            model_line.writeFitsData(self.__outPrefix + '.eline_model.fits')
            line_res.writeFitsData(self.__outPrefix + '.eline_res.fits')
        else:
            model_line.writeFitsData(self.__outPrefix + '.eline_model.fits',store_wave=True)
            line_res.writeFitsData(self.__outPrefix + '.eline_res.fits',store_wave=True)
        indices = numpy.arange(len(out_lines[2]))
        valid = numpy.zeros(len(out_lines[2]),dtype="bool")
        for i in range(len(out_lines[2])):
            if self.__datatype == 'CUBE':
                select_pos = (x_cor==out_lines[2][i]) & (y_cor==out_lines[3][i])
            elif self.__datatype == 'RSS':
                select_pos= fiber == out_lines[2][i]

            if numpy.sum(select_pos)>0:
                valid[i]=True
        columns = []
        if self.__datatype == 'CUBE':
            columns.append(pyfits.Column(name='x_cor', format='J', array=out_lines[2][valid]))
            columns.append(pyfits.Column(name='y_cor', format='J', array=out_lines[3][valid]))
        elif self.__datatype == 'RSS':
            columns.append(pyfits.Column(name='fiber', format='J', array=out_lines[2][valid]))
        for n in line_par._names:
            if line_par._profile_type[n] == 'Gauss':
                columns.append(pyfits.Column(name='%s_flux' % (n), format='E', array=out_lines[0][n]['flux'][valid]))
                columns.append(pyfits.Column(name='%s_vel' % (n), format='E', unit='km/s',
                    array=out_lines[0][n]['vel'][valid]))
                columns.append(pyfits.Column(name='%s_fwhm' % (n), format='E', unit='km/s',
                        array=out_lines[0][n]['fwhm'][valid]))

        try:
            table_out = pyfits.BinTableHDU.from_columns(columns)
        except:
            table_out = pyfits.new_table(columns)
        table_out.writeto(self.__outPrefix + '.eline_table.fits', overwrite=True)

    def run_bootstrap(self, stellar_parfile, eline_parfile, bootstraps, modkeep, parallel, verbose):
        """The bootstrap functions performs a bootstrap on the data in order to
        obtain errors. The errors for the template fitting are determined by
        refitting with a subset of the templates (with fixed velocity and
        velocity dispersion) and looking at the spread of the determined
        parameters to obtain a bootstrapped error estimate. Each time the
        templates are fitted, the emission lines are also fitted and this then
        gives an estimate of the error in the emission line parameters.

        Parameters
        ----------
        stellar_parfile : str
            The parameter file containing the constraints under which the input
            template spectra are fitted.
        eline_par_file : str
            The parameter file containing the constraints under which the
            emission lines are fitted to the input spectra.
        bootstraps : int
            The number of bootstraps run to each spectra.
        modkeep : float
            The percentage of template spectra that will be keeped under each
            bootstrap run.
        parallel : {'auto', int}, optional
            If parallel is not equal to one, the python multiprocessing routine
            shall be used to run parts of the code in parallel. With the option
            `auto`, it adjusts the number of parallel processes to the number
            of cpu-cores/threads that are available.
        verbose : bool, optional
            Produces screen output, such that the user can follow which spectrum
            is currently fitted and what the results are for each spectrum.

        Notes
        -----
        The function does not return anything, but overwrites the stellar table
        and emission-line table fits files created by the fitting functions
        `run_SSP_fit` and `run_eline_fit`. The files
        `outprefix`.stellar_table.fits and `outprefix`.eline_table.fits are
        appended with information on the errors on the fitted parameters.
        """
        ## Read in parameter file and get values
        parList = PyParadise.parameters.ParameterList(stellar_parfile)
        tmpldir = parList['tmpldir'].getValue()
        tmplfile = parList['tmplfile'].getValue()
        vel_guess = float(parList['vel_guess'].getValue())
        vel_min = float(parList['vel_min'].getValue())
        vel_max = float(parList['vel_max'].getValue())
        nwidth_norm = int(parList['nwidth_norm'].getValue())
        start_wave = float(parList['start_wave'].getValue())
        end_wave = float(parList['end_wave'].getValue())
        excl_fit = PyParadise.parameters.CustomMasks(parList['excl_fit'].getValue())
        excl_cont = PyParadise.parameters.CustomMasks(parList['excl_cont'].getValue())
        nlib_guess = int(parList['tmplinitspec'].getValue())
        kin_bootstrap = bool(int(parList['kin_bootstrap'].getValue()))

        bins = [[None, None]]
        try:
            f = open(parList['agebinfile'].getValue())
            try:
                for l in f.readlines():
                    bins.append([float(l.split()[0]), float(l.split()[1])])
            finally:
                f.close()
        except (IOError, ValueError):
            if verbose:
                print('agebinfile is non-existent or ignored.')
                print('Optional stellar population binning will not occur.')

        if nlib_guess < 0:
            modkeep = 100.0

        hdu = pyfits.open(self.__outPrefix + '.stellar_table.fits')
        stellar_table = hdu[1].data
        if self.__datatype == 'CUBE':
            x_cor = stellar_table.field('x_cor')
            y_cor = stellar_table.field('y_cor')
        elif self.__datatype == 'RSS':
            fiber = stellar_table.field('fiber')
        vel = stellar_table.field('vel_fit')
        disp = stellar_table.field('disp_fit')
        if kin_bootstrap:
            vel_err = stellar_table.field('vel_fit_err')
            disp_err = stellar_table.field('disp_fit_err')
        else:
            vel_err = None
            disp_err = None

        if eline_parfile is not None:
            parList = PyParadise.parameters.ParameterList(eline_parfile)
            eCompFile = parList['eCompFile'].getValue()
            vel_guess = float(parList['vel_guess'].getValue())
            line_fit = PyParadise.parameters.CustomMasks(parList['line_fit_region'].getValue())
            efit_method = parList['efit_method'].getValue()
            efit_ftol = float(parList['efit_ftol'].getValue())
            efit_xtol = float(parList['efit_xtol'].getValue())
            guess_window = int(parList['eguess_window'].getValue())

            line_par = PyParadise.fit_profile.parFile(eCompFile, self.__instrFWHM)

            hdu = pyfits.open(self.__outPrefix + '.eline_table.fits')
            eline_table = hdu[1].data
            if self.__datatype == 'CUBE':
                x_eline = eline_table.field('x_cor')
                y_eline = eline_table.field('y_cor')
            if self.__datatype == 'RSS':
                fiber_eline = eline_table.field('fiber')

        if verbose:
            print("The stellar population library is being prepared.")
        lib = PyParadise.ssplibrary.SSPlibrary(filename=tmpldir + '/' + tmplfile)
        min_wave = (start_wave / (1 + (vel_max + 2000) / 300000.0))
        max_wave = (end_wave / (1 + (vel_min - 2000) / 300000.0))

        if nlib_guess < 0:
            select = numpy.arange(lib.getBaseNumber()) == nlib_guess * -1 - 1
            lib = lib.subLibrary(select)
        lib = lib.subWaveLibrary(min_wave=min_wave, max_wave=max_wave)
        lib = lib.matchInstFWHM(self.__instrFWHM, vel_guess)
        lib = lib.resampleWaveStepLinear(self.__inputData.getWaveStep(), vel_guess / 300000.0)
        lib_norm = lib.normalizeBase(nwidth_norm, excl_cont, vel_guess / 300000.0)
        lib_rebin = lib_norm.rebinLogarithmic()

        if verbose:
            print("The input data is being normalized.")
        normData = self.__inputData.normalizeSpec(nwidth_norm, excl_cont.maskPixelsObserved(self.__inputData.getWave(),
             vel_guess / 300000.0))
        #normData.writeFitsData('test.fits')
        normDataSub = normData.subWaveLimits(start_wave, end_wave)
        excl_fit = excl_fit.maskPixelsObserved(normDataSub.getWave(), vel_guess / 300000.0)
        if eline_parfile is None:
            if self.__datatype == 'CUBE':
                (coeffs, maps, x_line, y_line) = normDataSub.fit_Lib_Boots(
                    lib_rebin, x_cor, y_cor, vel, disp, vel_err, disp_err,
                    mask_fit=excl_fit, bootstraps=bootstraps, modkeep=modkeep,
                    parallel=parallel, verbose=verbose)
            elif self.__datatype == 'RSS':
                (coeffs, maps) = normDataSub.fit_Lib_Boots(
                    lib_rebin, fiber, vel, disp, vel_err, disp_err,
                    mask_fit=excl_fit, bootstraps=bootstraps, modkeep=modkeep,
                    parallel=parallel, verbose=verbose)
        else:
            select_wave_eline = line_fit.maskPixelsObserved(normDataSub.getWave(), vel_guess / 300000.0)
            if self.__datatype == 'CUBE':
                (coeffs, maps, x_line, y_line) = normDataSub.fit_Lib_Boots(
                    lib_rebin, x_cor, y_cor, vel, disp, vel_err, disp_err,
                    line_par, select_wave_eline, excl_fit, efit_method,
                    guess_window, self.__instrFWHM, efit_ftol, efit_xtol,
                    bootstraps, modkeep, parallel, verbose)
            elif self.__datatype == 'RSS':
                (coeffs, maps) = \
                    normDataSub.fit_Lib_Boots(lib_rebin, fiber, vel, disp, vel_err, disp_err, bootstraps=bootstraps, par_eline=line_par,
                    select_wave_eline=select_wave_eline, mask_fit=excl_fit, method_eline=efit_method,
                    guess_window=guess_window, spectral_res=self.__instrFWHM, ftol=efit_ftol, xtol=efit_xtol,
                    modkeep=modkeep, parallel=parallel, verbose=verbose)
        mass_weighted_pars_full = numpy.zeros((len(vel), bootstraps, 5, len(bins) + 1), dtype=numpy.float32)
        lum_weighted_pars_full = numpy.zeros((len(vel), bootstraps, 5, len(bins) + 1), dtype=numpy.float32)
        for i in range(len(vel)):
            for j in range(len(bins)):
                for m in range(bootstraps):
                    try:
                        mass_weighted_pars_full[i, m, :, j] = lib_norm.massWeightedPars(coeffs[i, m, :], min_age=bins[j][0], max_age=bins[j][1])
                    except:
                        mass_weighted_pars_full[i, m, :, j] = numpy.array([numpy.nan, numpy.nan, numpy.nan, numpy.nan, numpy.nan])
                    try:
                        lum_weighted_pars_full[i, m, :, j] = lib_norm.lumWeightedPars(coeffs[i, m, :], min_age=bins[j][0], max_age=bins[j][1])
                    except:
                        lum_weighted_pars_full[i, m, :, j] = numpy.array([numpy.nan, numpy.nan, numpy.nan, numpy.nan, numpy.nan])
        mass_weighted_pars_mean = numpy.nanmean(mass_weighted_pars_full, axis=1)
        mass_weighted_pars_err = numpy.nanstd(mass_weighted_pars_full, axis=1)
        lum_weighted_pars_mean = numpy.nanmean(lum_weighted_pars_full, axis=1)
        lum_weighted_pars_err = numpy.nanstd(lum_weighted_pars_full, axis=1)
        columns_stellar = []
        for i, postfix in enumerate(['total'] + ['bin{}'.format(i) for i in range(1, len(bins))]):
            columns_stellar.append(pyfits.Column(name='lum_coeff_frac_' + postfix + '_btmean', format='E', array=lum_weighted_pars_mean[:, 0, i]))
            columns_stellar.append(pyfits.Column(name='lum_coeff_frac_' + postfix + '_err', format='E', array=lum_weighted_pars_err[:, 0, i]))
            columns_stellar.append(pyfits.Column(name='lum_age_' + postfix + '_btmean', format='E', array=lum_weighted_pars_mean[:, 1, i]))
            columns_stellar.append(pyfits.Column(name='lum_age_' + postfix + '_err', format='E', array=lum_weighted_pars_err[:, 1, i]))
            columns_stellar.append(pyfits.Column(name='lum_M/L_' + postfix + '_btmean', format='E', array=lum_weighted_pars_mean[:, 2, i]))
            columns_stellar.append(pyfits.Column(name='lum_M/L_' + postfix + '_err', format='E', array=lum_weighted_pars_err[:, 2, i]))
            columns_stellar.append(pyfits.Column(name='lum_[Fe/H]_' + postfix + '_btmean', format='E', array=lum_weighted_pars_mean[:, 3, i]))
            columns_stellar.append(pyfits.Column(name='lum_[Fe/H]_' + postfix + '_err', format='E', array=lum_weighted_pars_err[:, 3, i]))
            columns_stellar.append(pyfits.Column(name='lum_[A/Fe]_' + postfix + '_btmean', format='E', array=lum_weighted_pars_mean[:, 4, i]))
            columns_stellar.append(pyfits.Column(name='lum_[A/Fe]_' + postfix + '_err', format='E', array=lum_weighted_pars_err[:, 4, i]))
            columns_stellar.append(pyfits.Column(name='mass_coeff_frac_' + postfix + '_btmean', format='E', array=mass_weighted_pars_mean[:, 0, i]))
            columns_stellar.append(pyfits.Column(name='mass_coeff_frac_' + postfix + '_err', format='E', array=mass_weighted_pars_err[:, 0, i]))
            columns_stellar.append(pyfits.Column(name='mass_age_' + postfix + '_btmean', format='E', array=mass_weighted_pars_mean[:, 1, i]))
            columns_stellar.append(pyfits.Column(name='mass_age_' + postfix + '_err', format='E', array=mass_weighted_pars_err[:, 1, i]))
            columns_stellar.append(pyfits.Column(name='mass_M/L_' + postfix + '_btmean', format='E', array=mass_weighted_pars_mean[:, 2, i]))
            columns_stellar.append(pyfits.Column(name='mass_M/L_' + postfix + '_err', format='E', array=mass_weighted_pars_err[:, 2, i]))
            columns_stellar.append(pyfits.Column(name='mass_[Fe/H]_' + postfix + '_btmean', format='E', array=mass_weighted_pars_mean[:, 3, i]))
            columns_stellar.append(pyfits.Column(name='mass_[Fe/H]_' + postfix + '_err', format='E', array=mass_weighted_pars_err[:, 3, i]))
            columns_stellar.append(pyfits.Column(name='mass_[A/Fe]_' + postfix + '_btmean', format='E', array=mass_weighted_pars_mean[:, 4, i]))
            columns_stellar.append(pyfits.Column(name='mass_[A/Fe]_' + postfix + '_err', format='E', array=mass_weighted_pars_err[:, 4, i]))

        columns_bootstrap = None
        try:
            if bool(int(parList['bootstrap_verb'].getValue())) != 0:
                columns_bootstrap = []
                for m in range(bootstraps):
                    if lib.getBaseNumber() > 1:
                        columns_bootstrap.append(pyfits.Column(
                            name='bootstrap_coeff_{}'.format(m),
                            format='%dE' % (lib.getBaseNumber()),
                            array=coeffs[:, m]))
                    else:
                        columns_bootstrap.append(pyfits.Column(
                            name='bootstrap_coeff_{}'.format(m),
                            format='E', array=coeffs[:, m].flatten()))
        except KeyError:  # in case bootstrap_verb is not in the parameters-file
            pass

        tbl_size = 8 + 2 if self.__datatype == 'CUBE' else 8 + 1
        tbl_size += 10 * len(bins)
        try:
            hdu = pyfits.BinTableHDU.from_columns(stellar_table.columns[:tbl_size] + pyfits.ColDefs(columns_stellar))
            if columns_bootstrap is not None:
                btunit = pyfits.BinTableHDU.from_columns(pyfits.ColDefs(columns_bootstrap))
        except:
            hdu = pyfits.new_table(stellar_table.columns[:tbl_size] + pyfits.new_table(columns_stellar).columns)
            if columns_bootstrap is not None:
                btunit = pyfits.new_table(columns_bootstrap)
        if columns_bootstrap is None:
            hdu = pyfits.HDUList([pyfits.PrimaryHDU([]), hdu])
        else:
            hdu = pyfits.HDUList([pyfits.PrimaryHDU([]), hdu, btunit])
        hdu.writeto(self.__outPrefix + '.stellar_table.fits', overwrite=True)

        if self.__datatype == 'CUBE' and eline_parfile is not None:
            mapping=numpy.zeros(len(x_eline),dtype=numpy.int32)
            indices = numpy.arange(len(x_line))
            for i in range(len(x_eline)):
                select_pos = (x_line==x_eline[i]) & (y_line==y_eline[i])
                if numpy.sum(select_pos)==1:
                    mapping[i]=indices[select_pos]
                else:
                    mapping[i]=-1
        elif self.__datatype == 'RSS' and eline_parfile is not None:
            mapping=numpy.zeros(len(fiber_eline),dtype=numpy.int16)
            indices = numpy.arange(len(fiber))
            for i in range(len(fiber_eline)):
                select_pos = fiber==fiber_eline[i]
                if numpy.sum(select_pos)==1:
                    mapping[i]=indices[select_pos]
                else:
                    mapping[i]=-1

        if maps is not None:
            columns_eline = []
            for n in line_par._names:
                if line_par._profile_type[n] == 'Gauss':
                    columns_eline.append(pyfits.Column(name='%s_flux_err' % (n), format='E', array=maps[n]['flux_err'][mapping]))
                    columns_eline.append(pyfits.Column(name='%s_vel_err' % (n), format='E', unit='km/s',
                        array=maps[n]['vel_err'][mapping]))
                    columns_eline.append(pyfits.Column(name='%s_fwhm_err' % (n), format='E', unit='km/s',
                            array=maps[n]['fwhm_err'][mapping]))

            if self.__datatype == 'CUBE':
                add_column=2
            else:
                add_column=1
            try:
                hdu = pyfits.BinTableHDU.from_columns(eline_table.columns[:len(columns_eline) + add_column] + pyfits.ColDefs(columns_eline))
            except:
                hdu = pyfits.new_table(eline_table.columns[:len(columns_eline) + 2] + pyfits.new_table(columns_eline).columns)
            hdu.writeto(self.__outPrefix + '.eline_table.fits', overwrite=True)

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="""
Program to model the stellar population model from spectrosopic data stored either in RSS or datacube format.
Estimated parameters are velocity, velocity disperion, the best fit continuum model, and the star formation history.
Additionally the program allows to model emission lines in the residual spectra and infer the errors using a bootstrap
Monte Carlo simulation taking systematic uncertainties of the continuum model estimation into account.""",
formatter_class=argparse.ArgumentDefaultsHelpFormatter, prog='Paradise')
    parser.add_argument('--version', action='version', version='Paradise version %s' % (__version__))

    parser.add_argument("input", type=str, help="""File name of the input datacube or RSS file. Please have a look at the
        documentation for the correct format of each file.""")
    parser.add_argument("outprefix", type=str, help="""Prefix used for nameing all the output file names.""")
    parser.add_argument("instrFWHM", type=str, help="""Instrumental spectral resolution of the input spectra.""")
    parser.add_argument("--SSP_par", type=str, default=None, help="""File name of the parameter file that controls the fitting procedure""")
    parser.add_argument("--line_par", type=str, default=None, help="""File name of the parameter file that controls the fitting procedure""")
    parser.add_argument("--bootstraps", type=int, default=None, help="""Number of bootstraps iterations per spectrum""")
    parser.add_argument("--modkeep", type=float, default=80, help="""Fraction of random SSP models used for each bootstrap.""")
    parser.add_argument("--parallel", type=str, default="auto", help="""Options are: 'auto' - using all CPUs available on the
    machine, an integer number specifying in the number of CPUs. An integer of 1 means no parrell processing.""")
    parser.add_argument("--verbose", action="store_true", default=False, help="""Flag to print some progress information on
    the screen. The default is False.""")

    args = parser.parse_args()
    app = ParadiseApp(args.input, args.outprefix, args.instrFWHM)
    if args.SSP_par is not None and args.bootstraps is None:
        app.run_SSP_fit(args.SSP_par, args.parallel, args.verbose)
    if args.line_par is not None and args.bootstraps is None:
        app.run_eline_fit(args.line_par, args.parallel, args.verbose)
    if args.bootstraps is not None:
        app.run_bootstrap(args.SSP_par, args.line_par, args.bootstraps, args.modkeep, args.parallel, args.verbose)


import numpy as np
import gc

from multiprocessing.pool import ThreadPool

from scipy.linalg import pinv
from scipy.sparse import csr_matrix
from scipy.sparse import lil_matrix

from .util import _checkState, _thresholding, _save, _load
from .distributed import *

from pyspark import keyword_only, SparkContext
from pyspark.ml import Estimator, Model
from pyspark.ml.evaluation import Evaluator
from pyspark.ml.param import *
from pyspark.ml.param.shared import HasMaxIter, HasStandardization, HasTol
from pyspark.ml.tuning import CrossValidator, CrossValidatorModel

class CCA(Estimator, HasMaxIter, HasStandardization, HasTol):
    verbose = Param(Params._dummy(), "verbose", "if True, print progress to STDOUT")
    k = Param(Params._dummy(), "k", "the number of canonical vectors",
              typeConverter=TypeConverters.toInt)

    rhos = Param(Params._dummy(), "rhos", "the sparsity parameters in [0,1]",
                 typeConverter=TypeConverters.toList)

    colsPerBlock = Param(Params._dummy(), "colsPerBlock",
                         "Number of columns in the blocks of the two input "
                         "matrices when they are distributed",
                         typeConverter=TypeConverters.toList)

    broadcast = Param(Params._dummy(), "broadcast",
                      "whether or not to broadcast during distributed matrix "
                      "multiplication",
                      typeConverter=TypeConverters.toBoolean)

    caching = Param(Params._dummy(), "caching",
                      "whether or not to cache underlying RDDs of ColBlockMatrices",
                      typeConverter=TypeConverters.toBoolean)

    @keyword_only
    def __init__(self, k=1, maxIter=1000, rhos=[0.1], colsPerBlock=[1024],
                 broadcast=True, caching=True, standardization=True, tol=1e-4,
                 verbose=False):
        super(CCA, self).__init__()
        self._setDefault(
            k=1, maxIter=1000, rhos=[0.1], colsPerBlock=[1024], broadcast=True,
            caching=True, standardization=True, tol=1e-4, verbose=False
        )
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    def setVerbose(self, verbose):
        """
        Sets the value of :py:attr:`verbose`.
        """
        return self._set(verbose=verbose)

    def getVerbose(self):
        """
        Gets the value of verbose or its default value.
        """
        return self.getOrDefault(self.verbose)

    @keyword_only
    def setParams(self, k=1, maxIter=1000, rhos=[0.1],
                  colsPerBlock=[1024], broadcast=True, caching=True,
                  standardization=True, tol=1e-4, verbose=False):
        self.setK(k)
        self.setMaxIter(maxIter)
        self.setRhos(rhos)
        self.setColsPerBlock(colsPerBlock)
        self.setBroadcast(broadcast)
        self.setCaching(caching)
        self.setStandardization(standardization)
        self.setTol(tol)
        self.setVerbose(verbose)

    def setRhos(self, value):
        """
        Sets the value of :py:attr:`rhos`.
        """
        assert min(value) >= 0 and max(value) <= 1, \
            "all rhos must be in [0, 1]"
        return self._set(rhos=value)

    def getRhos(self):
        """
        Gets the value of :py:attr:`rhos` or its default value.
        """
        return self.getOrDefault(self.rhos)

    def setColsPerBlock(self, value):
        """
        Sets the value of :py:attr:`colsPerBlock`.
        """
        # assert len(value) == 2, "colsPerBlock must be a two-tuple"
        assert min(value) > 0, "colsPerBlock must be positive"
        return self._set(colsPerBlock=value)

    def getColsPerBlock(self):
        """
        Gets the value of :py:attr:`colsPerBlock`. or its default value.
        """
        return self.getOrDefault(self.colsPerBlock)

    def setBroadcast(self, value):
        """
        Sets the value of :py:attr:`broadcast`.
        """
        return self._set(broadcast=value)

    def getBroadcast(self):
        """
        Gets the value of :py:attr:`broadcast` or its default value.
        """
        return self.getOrDefault(self.broadcast)

    def setCaching(self, value):
        """
        Sets the value of :py:attr:`caching`.
        """
        return self._set(caching=value)

    def getCaching(self):
        """
        Gets the value of :py:attr:`caching` or its default value.
        """
        return self.getOrDefault(self.caching)

    def setK(self, value):
        """
        Sets the value of :py:attr:`k`.
        """
        assert value > 0, "k must be greater than 0"
        return self._set(k=value)

    def getK(self):
        """
        Gets the value of k or its default value.
        """
        return self.getOrDefault(self.k)

    def _fit(self, datasets):
        """
        Computes multiple CCs using a deflation scheme.
        :param datasets: A list of numpy.arrays.
        """
        standardization = self.getStandardization()
        n_ccs = self.getK()
        maxIter = self.getMaxIter()
        rhos = self.getRhos()
        tol = self.getTol()
        colsPerBlock = self.getColsPerBlock()
        broadcast = self.getBroadcast()
        caching = self.getCaching()
        verbose = self.getVerbose()

        if standardization:
            datasets = [X - np.mean(X, axis=0) for X in datasets]

        d = len(datasets)

        if len(rhos) == 1:
            rhos = rhos * d
        else:
            assert len(rhos) == d,\
                "Please provide one regularization parameter per dataset."
        if len(colsPerBlock) == 1:
            colsPerBlock = colsPerBlock * d
        else:
            assert len(colsPerBlock) == d,\
                "Please provide one colsPerBlock value per dataset."

        # cross covariance matrices
        CC = [
            makeColBlockMatrix(datasets[i], colsPerBlock[i], caching)
            .leftMultiply(datasets[j].T, broadcast)
            for i in range(d)
            for j in (x for y in (range(i), range(i+1, d))
                      for x in y)  # all indices other than i
        ]

        # matrices of sparse loadings
        ZZ = [lil_matrix((n_ccs, X.shape[1])) for X in datasets]
        Sigma = []

        for k in range(n_ccs):

            z_init = []  # initial estimates of sparse canonical loadings
            pattern = []  # sparsity patterns

            for i in range(d):

                # 1. Get the cov matrices involving dataset i. For any
                #    covariance matrix involving a dataset that was
                #    already processed, reduce the rows using the sparsity
                #    pattern that was computed.
                # 2. Call singleL1 one time for all of those cov matrices,
                #    returning a z for each one.
                # 3. Add the z's up from step 2 and threshold and standardize
                #    the resulting vector. This is z initial for dataset i.

                z = np.zeros(datasets[i].shape[1])
                max_rho = 0

                # indices of datasets other than i
                not_i = [x for y in (range(i), range(i+1, d)) for x in y]

                # indices of covariance matrices we'll need for this iteration
                for j in range(i*(d - 1), (i + 1)*(d - 1)):

                    # get the index of the left dataset in CC[j]
                    left_idx = not_i[j % (d - 1)]

                    if left_idx < i:
                        p = pattern[left_idx]
                        if np.size(p) == 0:
                            print("Canonical loading vector " + str(k) +
                                  " of dataset " + str(left_idx) +
                                  " is all zeros. Moving on to next dataset.")
                            # move on to next not_i dataset
                            import pdb; pdb.set_trace()
                            continue
                    else:
                        p = None  # keep all columns of left dataset

                    if k == 0:
                        lZ = None
                        rZ = None
                    else:
                        lZ = (ZZ[left_idx][0:k, p] if p is not None
                              else ZZ[left_idx][0:k, :]).tocsr()
                        sigma_ij = np.diag(
                            [Sigma[cc][i, left_idx] for cc in range(k)]
                        )
                        rZ = csr_matrix.dot(
                            sigma_ij,
                            ZZ[i][0:k, :].tocsr()
                        ) ## TODO: check

                    z_term, this_rho = CCA._singleL1(
                        CC[j], datasets[left_idx][:, p] if p is not None
                        else datasets[left_idx], datasets[i],
                        rhos[i], tol, maxIter, lZ, rZ, p, verbose
                    )
                    z += z_term
                    if this_rho > max_rho:
                        max_rho = this_rho

                z = _thresholding(z, max_rho)
                norm_z = np.linalg.norm(z)
                if norm_z > 0:
                    z = z / norm_z

                pattern.append(z.nonzero()[0])
                z_init.append(z)

            zz, sigma = CCA._coeffEstSingleComp(
                datasets, z_init, pattern, k, maxIter, tol, ZZ, Sigma, verbose
            )
            for i in range(d):
                ZZ[i][k, pattern[i]] = zz[i]
            Sigma.append(sigma)
            Sigma[k] += np.tril(Sigma[k]).T

            for i in range(d):
                for j in (x for y in (range(i), range(i+1, d)) for x in y):
                    # C_ij - z_i.T C_ij z_j * z_i z_j.T
                    # subtractOuter mutates its object, so no need to reassign to CC
                    CC[i*(d - 1) + (j if j < i else j - 1)].subtractOuter(
                        Sigma[k][i, j] * ZZ[j][k, :].todense().getA1(),
                        ZZ[i][k, :].todense().getA1()
                    )
        ## done, unpersist the ColBlockMatrices
        if caching:
            for cbm in CC:
                cbm.unpersist()
        del CC
        gc.collect()
        return CCAModel(k=n_ccs, ZZ=[Z.tocsr() for Z in ZZ], Sigma=Sigma,
                        standardization=standardization, rhos=rhos)

    def fit(self, datasets, params=None):
        """
        Overrides pyspark.ml.Estimator's fit method.
        """
        # TODO: assert all datasets are np.array
        assert len(datasets) > 1, "need at least 2 datasets"
        assert np.all(
            np.array([X.shape[0] for X in datasets]) == datasets[0].shape[0]),\
            "all datasets must have same number of observations"
        if params is None:
            params = dict()
        if isinstance(params, (list, tuple)):
            models = [None] * len(params)
            for index, model in self.fitMultiple(datasets, params):
                models[index] = model
            return models
        elif isinstance(params, dict):
            if params:
                return self.copy(params)._fit(datasets)
            else:
                return self._fit(datasets)
        else:
            raise ValueError("Params must be either a param map or a list/"
                             "tuple of param maps, but got %s." % type(params))

    @staticmethod
    def _coeffEstSingleComp(datasets, z_init, pattern, k, maxIter, tol,
                            ZZ=None, Sigma=None, verbose=False):
        """
        """
        d = len(datasets)
        z_final = []

        for i in range(d):

            # get reduced versions of the ith dataset and initial loadings
            # vector using the ith sparsity pattern
            p_i = pattern[i]
            if p_i.size == 0:
                continue
            X_i = datasets[i][:, p_i]
            z_i = z_init[i][p_i]

            nIter = 1
            f = []

            while not _checkState(f, nIter, maxIter, tol, verbose) in (-1, 1, 2):

                tmp = np.zeros(p_i.size)

                # all indices other than i
                for j in (x for y in (range(i), range(i+1, d)) for x in y):

                    p_j = pattern[j]
                    if np.size(p_j) == 0:
                        continue
                    X_j = datasets[j][:, p_j]
                    if j < i:
                        z_j = z_final[j]  # already shrunken
                        tmp += np.matmul(  # X_i.T @ X_j @ z_j
                            X_i.T, np.matmul(
                                X_j, z_j
                            )
                        )
                        if k > 0:
                            tmp -= Sigma[k-1][i, j] * csr_matrix.dot(
                                ZZ[j][k-1, p_j].tocsr(), z_j
                            ) * ZZ[i][k-1, p_i].tocsr().transpose().\
                                todense().getA1()
                    else:
                        tmp += np.matmul(  # X_i.T @ X_j @ X_j.T @ X_i @ z_i
                            X_i.T, np.matmul(
                                X_j, np.matmul(
                                    X_j.T, np.matmul(X_i, z_i)
                                )
                            )
                        )
                        if k > 0:
                            sig = Sigma[k-1][i, j]
                            zz_i = ZZ[i][k-1, p_i].tocsr()  # row vector
                            zz_j = ZZ[j][k-1, p_j].tocsr()  # row vector
                            tmp -= sig * np.dot(
                                csr_matrix.dot(
                                    zz_j, X_j.T
                                ),
                                np.matmul(X_i, z_i)
                            ) * zz_i.transpose().todense().getA1() +\
                                sig * csr_matrix.dot(zz_i, z_i) *\
                                np.matmul(
                                    X_i.T, csr_matrix.dot(
                                        X_j, zz_j.transpose()
                                    )
                                ).ravel() - sig**2 * csr_matrix.dot(
                                    zz_j, zz_j.transpose()
                                ).todense().getA1() * csr_matrix.dot(
                                    zz_i, z_i
                                ) * zz_i.transpose().todense().getA1()

                norm_tmp = np.linalg.norm(tmp)
                f.append(-2*norm_tmp)
                if norm_tmp > 0:
                    z_i = tmp / norm_tmp
                else:
                    z_i = tmp
                nIter += 1

            z_final.append(z_i.ravel())

        sigma = np.zeros((d, d))
        for i in range(d):
            p_i = pattern[i]
            X_i = datasets[i][:, p_i]
            z_i = z_final[i]
            if k > 0:
                z_i_prev = ZZ[i][k-1, p_i].todense().getA1()  # row vector
            for j in range(i):
                p_j = pattern[j]
                if p_i.size == 0 or p_j.size == 0:
                    sigma[i, j] = 0
                    continue
                X_j = datasets[j][:, p_j]
                z_j = z_final[j]
                if k == 0:
                    # z_i.T @ X_i.T X_j z_j for j < i
                    sigma[i, j] = np.dot(
                        np.matmul(X_i, z_i),
                        np.matmul(X_j, z_j)
                    )
                else:
                    z_j_prev = ZZ[j][k-1, p_j].todense().getA1()
                    # z_i.T @ (X_i.T @ X_j - Sigma[k-1][i, j] *
                    #     ZZ[i][k-1,p_i] @ ZZ[j][k-1,p_j].T) @ z_j
                    sigma[i, j] = np.dot(
                        np.matmul(X_i, z_i),
                        np.matmul(X_j, z_j)
                    ) - Sigma[k-1][i, j] *\
                        np.dot(z_i, z_i_prev) * np.dot(z_j, z_j_prev)

        return (z_final, sigma)

    @staticmethod
    def _singleL1(covBlock, lX, rX, rho, tol, maxIter,
                  lZ=None, rZ=None, p=None, verbose=False):
        """
        Using an off-diagonal block of the sample covariance matrix,
        `covBlock`, return an initial estimate of the canonical loading
        vector corresponding to the `right` matrix of `covBlock`.

        :param covBlock: A sparkle.distributed.ColBlockMatrix formed by
                       multiplying the transpose of the `lX` matrix by
                       the `rX` matrix.
        :param lX: The left matrix in the sample covariance block.
        :param rX: The right matrix in the sample covariance block.
        :param lZ: A scipy.sparse.csr_matrix which contains the loadings
                   corresponding to the left matrix (optional).
        :param rZ: A scipy.sparse.csr_matrix which contains the loadings
                   corresponding to the right matrix (optional).
        :param rho: Sparsity parameter.
        :param tol: Convergence tolerance.
        :param maxIter: Maxiumum number of iterations.
        :param p: Optional sparsity pattern.
        """
        rho_max, x = covBlock.maxL2NormCol(p)
        x = x / rho_max
        rho = rho * rho_max

        f = []
        nIter = 1

        while not _checkState(f, nIter, maxIter, tol, verbose):

            z = np.matmul(rX.T, np.matmul(lX, x))
            if lZ is not None:
                z = z - rZ.transpose().dot(lZ.dot(x))
            z = _thresholding(z, rho)

            f.append(np.sum(z**2))

            x = np.matmul(lX.T, np.matmul(rX, z))
            if lZ is not None:
                x = x - lZ.transpose().dot(rZ.dot(z))
            norm_x = np.linalg.norm(x)
            if norm_x > 0:
                x = x / norm_x

            nIter += 1

        z = np.matmul(rX.T, np.matmul(lX, x))
        if lZ is not None:
            z = z - rZ.transpose().dot(lZ.dot(x))
        return z, rho

    def suggest_colsPerBlock(datasets, n_cores=4, partitions_per_core=3):
        """
        :param partitions_per_core: Spark recommends 2-4 partitions per CPU.
          See https://spark.apache.org/docs/latest/rdd-programming-guide.html
        """
        nCols = np.array([X.shape[1] for X in datasets])
        nPartitions = n_cores * partitions_per_core
        return (nCols / nPartitions).astype(int)

class CCAModel(Model, HasStandardization):

    k = Param(Params._dummy(), "k", "the number of canonical vectors",
              typeConverter=TypeConverters.toInt)

    rhos = Param(Params._dummy(), "rhos", "the sparsity parameters in [0,1]",
                 typeConverter=TypeConverters.toList)

    @keyword_only
    def __init__(self, k=1, ZZ=None, Sigma=None, standardization=True,
                 rhos=[0.1]):
        super(CCAModel, self).__init__()
        self._setDefault(k=1, standardization=True, rhos=[0.1])
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    @keyword_only
    def setParams(self, k=1, ZZ=None, Sigma=None, standardization=True,
                  rhos=[0.1]):
        self.setK(k)
        self.setStandardization(standardization)
        self.setRhos(rhos)
        self.__ZZ = ZZ
        self.__Sigma = Sigma

    def setK(self, value):
        """
        Sets the value of :py:attr:`k`.
        """
        assert value > 0, "k must be greater than 0"
        return self._set(k=value)

    def getK(self):
        """
        Gets the value of k or its default value.
        """
        return self.getOrDefault(self.k)

    def setRhos(self, value):
        """
        Sets the value of :py:attr:`rhos`.
        """
        assert min(value) >= 0 and max(value) <= 1,\
            "all rhos must be in [0, 1]"
        return self._set(rhos=value)

    def getRhos(self):
        """
        Gets the value of :py:attr:`rhos` or its default value.
        """
        return self.getOrDefault(self.rhos)

    @property
    def ZZ(self):
        """
        A np.array of scipy.sparse.csr_matrix with first :py:attr:`k` canonical
        weight vectors.
        """
        return self.__ZZ

    @property
    def Sigma(self):
        """
        A np.array of scipy.sparse.csr_matrix with first :py:attr:`k` canonical
        weights on the diagonal.
        """
        return self.__Sigma

    def canonicalCorrelations(self, datasets = None):
        """
        Take :py:attr:`k` datasets and return the :py:attr:`k` canonical correlations.
        :param datasets: A list of numpy.arrays.
        """
        if datasets is not None:
            XZ = [csr_matrix.dot(X, Z.transpose()) for X, Z in
                  zip(datasets, self.ZZ)]
            k = self.getK()
            corrs = np.zeros((k, len(XZ), len(XZ)))
            for i in range(len(XZ)):
                corrs[:, i, i] = np.ones(k)
                for j in range(i + 1, len(XZ)):
                    for cc in range(k):
                        corrs[cc, j, i] = np.corrcoef(
                            XZ[i][:, cc], XZ[j][:, cc], rowvar=False)[0, 1]
            for cc in range(k):
                # fill in the upper triangles with the transpose of the lower
                corrs[cc, :, :] = corrs[cc, :, :] + np.tril(
                    corrs[cc, :, :], -1).T
        else:
            corrs = None  # TODO: save canonical correlations during CCA.fit
        return np.nan_to_num(corrs)

    def _transform(self, datasets, outcome_index=None):
        # TODO: use 'outcome_index' to allow user to use d - 1 datasets to predict the
        # remaining one
        ZZ = self.ZZ
        d = len(ZZ)
        assert len(datasets) == d,\
            "number of datasets should be len(self.ZZ)"
        standardization = self.getStandardization()
        if standardization:
            datasets = [X - np.mean(X, axis=0) for X in datasets]

        if outcome_index is not None:
            assert outcome_index >= 0 and outcome_index < d,\
                "outcome_index is not a valid index for datasets"
            XZ = [csr_matrix.dot(datasets[j], ZZ[j].transpose()) for j in
                  [x for x in range(d) if x != outcome_index]]
            XZ_sums = np.sum(XZ, axis=0)
            prediction = np.dot(XZ_sums, pinv(ZZ[outcome_index].todense()).transpose())
        else:
            XZ = [csr_matrix.dot(X, Z.transpose()) for X, Z in zip(datasets, ZZ)]
            XZ_sums = []
            for i in range(d):
                for j in range(d):
                    if j != i:
                        if len(XZ_sums) == i:
                            XZ_sums.append(XZ[j])
                        else:
                            XZ_sums[i] += XZ[j]
            prediction = [np.dot(XZ_sums[i], pinv(ZZ[i].todense()).transpose())
                          for i in range(d)]
        return prediction

    def transform(self, datasets, params=None, outcome_index=None):
        """
        Overrides pyspark.ml.Transformer's transform method.
        :param datasets: A list of np.arrays corresponding to self.ZZ
        :param params: an optional param map that overrides embedded params.
        :returns: predictions
        """
        if params is None:
            params = dict()
        if isinstance(params, dict):
            if params:
                return self.copy(params)._transform(datasets, outcome_index)
            else:
                return self._transform(datasets, outcome_index)
        else:
            raise ValueError("Params must be a param map but got %s." % type(params))

    def save(self, path):
        """Save this CCAModel instance to the given path"""
        _save(self, path)

    def load(path):
        """Load a CCAModel instance from the given path"""
        return _load(path)

class CCACanonicalCorrelationEvaluator(Evaluator):

    def _evaluate(self, canCorrs):
        # TODO: extract upper triangles
        return np.mean(canCorrs)

    def evaluate(self, model, datasets, params=None):
        """
        Overrides pyspark.ml.evaluation.Evaluator's evaluate method.
        :param canCorrs: the result of a call to CCAModel.canonicalCorrelations
        :returns: average correlation between the canonical variables
        """
        if params is None:
            params = dict()
        if isinstance(params, dict):
            if params:
                return self.copy(params)._evaluate(
                    model.canonicalCorrelations(datasets)
                )
            else:
                return self._evaluate(
                    model.canonicalCorrelations(datasets)
                )
        else:
            raise ValueError("Params must be a param map but got %s."
                             % type(params))

class CCAPredictionEvaluator(Evaluator):

    outcome_index = Param(Params._dummy(), "outcome_index",
                          "the index of the outcome dataset",
                          typeConverter=TypeConverters.toInt)

    @keyword_only
    def __init__(self, outcome_index=0):
        super(Evaluator, self).__init__()
        self._setDefault(outcome_index=0)
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    @keyword_only
    def setParams(self, outcome_index=0):
        self.setOutcomeIndex(outcome_index)

    def setOutcomeIndex(self, value):
        """
        Sets the value of :py:attr:`outcome_index`.
        """
        return self._set(outcome_index=value)

    def getOutcomeIndex(self):
        """
        Gets the value of :py:attr:`outcome_index`.
        """
        return self.getOrDefault(self.outcome_index)

    def _evaluate(self, prediction, dataset):
        corrs = np.zeros(prediction.shape[1])
        for i in range(prediction.shape[1]):
            corrs[i] = np.corrcoef(prediction[:, i], dataset[:, i])[0, 1]
        return np.mean(np.nan_to_num(corrs))

    def evaluate(self, model, datasets, params=None):
        """
        Overrides pyspark.ml.evaluation.Evaluator's evaluate method.
        :param canCorrs: the result of a call to CCAModel.canonicalCorrelations
        :returns: average correlation between the canonical variables
        """
        if params is None:
            params = dict()
        if isinstance(params, dict):
            if params:
                return self.copy(params)._evaluate(
                    model.transform(datasets, outcome_index=self.getOutcomeIndex()),
                    datasets[self.getOutcomeIndex()]
                )
            else:
                return self._evaluate(
                    model.transform(datasets, outcome_index=self.getOutcomeIndex()),
                    datasets[self.getOutcomeIndex()]
                )
        else:
            raise ValueError("Params must be a param map but got %s."
                             % type(params))

class CCACrossValidator(CrossValidator):
    """
    Borrows heavily from pyspark.ml.tuning.CrossValidator.
    """
    verbose = Param(Params._dummy(), "verbose", "if True, print progress to STDOUT")

    @keyword_only
    def __init__(self, estimator=None, estimatorParamMaps=None, evaluator=None,
                 numFolds=3, seed=None, parallelism=1, verbose=False):
        super(CCACrossValidator, self).__init__(
            estimator=estimator,
            estimatorParamMaps=estimatorParamMaps,
            evaluator=evaluator,
            numFolds=numFolds, seed=seed,
            parallelism=parallelism
        )
        self._setDefault(verbose=False)
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    @keyword_only
    def setParams(self, estimator=None, estimatorParamMaps=None, evaluator=None,
                  numFolds=3, seed=None, parallelism=1, verbose=False):
        self.setEstimator(estimator)
        self.setEstimatorParamMaps(estimatorParamMaps)
        self.setEvaluator(evaluator)
        self.setNumFolds(numFolds)
        self.setSeed(seed)
        self.setParallelism(parallelism)
        self.setVerbose(verbose)

    def setVerbose(self, verbose):
        """
        Sets the value of :py:attr:`verbose`.
        """
        return self._set(verbose=verbose)

    def getVerbose(self):
        """
        Gets the value of verbose or its default value.
        """
        return self.getOrDefault(self.verbose)

    def _fit(self, datasets):
        """
        Overrides pyspark.ml.tuning.Crossvalidator's _fit method.
        """
        verbose = self.getVerbose()
        est = self.getOrDefault(self.estimator)
        epm = self.getOrDefault(self.estimatorParamMaps)
        numModels = len(epm)
        eva = self.getOrDefault(self.evaluator)
        nFolds = self.getOrDefault(self.numFolds)
        h = 1.0 / nFolds
        randUnif = np.random.uniform(size = datasets[0].shape[0])
        metrics = np.zeros(numModels)

        pool = ThreadPool(processes=min(self.getParallelism(), numModels))

        for i in range(nFolds):
            if verbose:
                print("Fold " + str(i + 1) + " of " + str(nFolds))
            validateLB = i * h
            validateUB = (i + 1) * h
            condition = (randUnif >= validateLB) & (randUnif < validateUB)
            validate = [X[condition,:] for X in datasets]
            train = [X[~condition,:] for X in datasets]

            tasks = self._parallelFitTasks(est, train, eva, validate, epm, verbose)
            for j, metric in pool.imap_unordered(lambda f: f(), tasks):
                metrics[j] += (metric / nFolds)

        if eva.isLargerBetter():
            bestIndex = np.argmax(metrics)
        else:
            bestIndex = np.argmin(metrics)
        bestModel = est.fit(datasets, epm[bestIndex])
        if verbose:
            print("Best model: " + str(i + 1) + " of " + str(nFolds))
            print("k: {0}".format(bestModel.getK()))
            print("rhos: {0}".format(bestModel.getRhos()))
            print("Avg. correlation across " + str(nFolds) + " folds: " +
                  str(metrics[bestIndex]))
        return CCACrossValidatorModel(bestModel)

    def fit(self, datasets, params=None):
        """
        Overrides pyspark.ml.Estimator's fit method, which is inherited from
        pyspark.ml.tuning.CrossValidator.
        """
        if params is None:
            params = dict()
        if isinstance(params, (list, tuple)):
            models = [None] * len(params)
            for index, model in self.fitMultiple(datasets, params):
                models[index] = model
            return models
        elif isinstance(params, dict):
            if params:
                return self.copy(params)._fit(datasets)
            else:
                return self._fit(datasets)
        else:
            raise ValueError("Params must be either a param map or a list/tuple of"
                             "param maps, but got %s." % type(params))

    @staticmethod
    def _parallelFitTasks(est, train, eva, validation, epm, verbose=False):
        """
        Creates a list of callables which can be called from different threads to fit and evaluate
        an estimator in parallel. Each callable returns an `(index, metric)` pair. Borrows heavily
        from pyspark.ml.tuning._parallelFitTasks.

        :param est: Estimator, the estimator to be fit.
        :param train: DataFrame, training data set, used for fitting.
        :param eva: Evaluator, used to compute `metric`
        :param validation: DataFrame, validation data set, used for evaluation.
        :param epm: Sequence of ParamMap, params maps to be used during fitting & evaluation.
        :return: (int, float), an index into `epm` and the associated metric value.
        """
        modelIter = est.fitMultiple(train, epm)

        def singleTask():
            index, model = next(modelIter)
            metric = eva.evaluate(model, validation)
            if verbose:
                print("k: {0}".format(model.getK()))
                print("rhos: {0}".format(model.getRhos()))
                print("Avg. correlation: " + str(metric))
            return index, metric

        return [singleTask] * len(epm)

class CCACrossValidatorModel(CrossValidatorModel):
    """
    Borrows heavily from pyspark.ml.tuning.CrossValidatorModel.
    """
    def __init__(self, bestModel):
        super(CCACrossValidatorModel, self).__init__(bestModel)

    def _transform(self, datasets):
        return self.bestModel.transform(datasets)

    def transform(self, datasets, params=None):
        """
        Overrides pyspark.ml.Transformer's transform method.
        """
        if params is None:
            params = dict()
        if isinstance(params, dict):
            if params:
                return self.copy(params)._transform(datasets)
            else:
                return self._transform(datasets)
        else:
            raise ValueError("Params must be a param map but got %s." % type(params))

    def save(self, path):
        """Save this CCACrossValidatorModel instance to the given path"""
        _save(self, path)

    def load(path):
        """Load a CCACrossValidatorModel instance from the given path"""
        return _load(path)


#!/usr/bin/env python

import os,sys,re
from ase import Atom, Atoms
import io2
import io2.data as data
from io2.gaussian_reader import GaussianReader as GR0
import ase.data as ad
import numpy as np
import multiprocessing

global dic_sm, atom_symbs, prop_names
dic_sm = {'H':2,'B':2,'C':3,'N':4,'O':3,'F':2, \
          'Si':3,'P':4,'S':3,'Cl':2,'Ge':3,'As':4,\
          'Se':3,'Br':2,'I':2}
atom_symbs = ['C', 'B', 'F', 'I', 'H', 'O', 'N', 'P', \
                  'S', 'Se', 'As', 'Cl', 'Si', 'Ge', 'Br']
prop_names = ['NMR', 'ALPHA', 'MU', 'AE', 'G','U0','U','H', \
              'HOMO','LUMO','GAP', 'IP','EA', 'OMEGA1', \
              'MP','BP']

def is_job_done(f):
    el = file(f).readlines()[-1] # end line
    return 'Normal termination' in el

class GR(object):

    def __init__(self, f, properties=['AE'], unit='kcal', istart=0):

        self.f = f
        self.unit = unit

        assert is_job_done(f), '#ERROR: Gaussian calculation was terminated abnormally?'

        dic = GR0(f, istart=istart)[-1]
        self.coords = dic['Positions']
        self.zs = dic['Atomic_numbers']
        self.atoms = Atoms(self.zs, self.coords)
        self.na = len(self.zs)
        self.sm = dic['Multiplicity']
        self.charge = dic['Charge']
        self.dic = dic

        uo = io2.Units()
        self.uo = uo
        self.const = {'kcal': uo.h2kc, 'kj': uo.h2kj, 'h': 1.0, \
                 'ev': uo.h2e, }[unit.lower()]

        if type(properties) is str: properties = [ properties.upper(), ]

        # visited or not?
        imo, ie, ithermo, imu, ipolar, inmr, iforce, ipop = 0, 0, 0, 0, 0, 0, 0, 0
        ys = []
        visited = []
        for property_i in properties:
            pname = property_i.upper()
            if pname in ['HOMO','LUMO','GAP']:
                if not imo:
                    imo = 1
                    self.get_mo_energy()
            elif pname in ['E', 'AE', 'EE', 'EN' ]:
                if not ie:
                    ie = 1
                    self.get_energy(etype=pname)
            elif pname in ['HF', 'MP2', 'CCSD(T)', ]:
                if pname not in visited:
                    self.get_ei(pname)
                    visited.append( pname )
            elif pname in ['U0', 'U', 'H', 'G', 'ZPE', 'ZPVE']:
                if not ithermo:
                    ithermo = 1
                    self.get_thermo()
            elif pname in ['MU',]:
                if not imu:
                    imu = 1
                    self.get_dipole_moment()
            elif pname in ['ALPHA',]:
                if not ipolar:
                    ipolar = 1
                    self.get_polarizability()
            elif pname in ['NMR',]:
                if not inmr:
                    inmr = 1
                    self.get_nmr()
            elif pname in ['FORCE',]:
                if not iforce:
                    iforce = 1
                    self.get_force()
            elif pname in ['POPULATION',]:
                if not ipop:
                    ipop = 1
                    self.get_mulliken_population()
            else:
                raise '#ERROR:'
            ys.append( self.dic[ property_i ] )
	    self.properties = ys

    def update_dic(self, ss1, ss2):
        for i, s1 in enumerate(ss1):
            self.dic[s1] = ss2[i]

    def get_ei(self, ein): # `ein -- energy_i name, i.e., HF, MP2, CCSD(T)
        self.update_dic([ein, ], [ self.dic[ein]*self.const, ])

    def get_energy(self,etype='E'):
        f = self.f
        # first read method
        cmd1 = "grep -n ' # ' %s | head -n 1"%f
        s0 = os.popen(cmd1).read().strip().split(':')
        #print 'f,s0 = ',f,s0;
        n0 = int(s0[0]); s0_1 = s0[1]

        # in case the command line has two many user-specified keywords
        # so that it spans more than 1 line
        s0u = s0_1
        while True:
            cmd1_i = "sed -n '%sp' %s"%(n0+1,f)
            s0_i = os.popen(cmd1_i).read().strip()
            if set([ si for si in s0_i ]) != set(['-',]):
                s0u += s0_i; n0 += 1
            else:
                break

        s1 = s0u.split()
        #print s1
        ns1 = len(s1)
        method1 = None
        isGn = False
        basis = ''
        # for MY calculation, the command line always goes like "# G4MP2"
        if ns1 == 2 and (s1[1][0].upper() == 'G'):
            isGn = True
            method1 = s1[1]
        else:
            # assume this is a normal method with "/" seperating `qcl and
            # `basis, e.g., B3LYP/6-31G(d,p)'
            for sj in s1:
                if '/' in sj:
                    #print 'Yeah, sj = ',sj
                    cont1 = sj.split('/')
                    method1 = cont1[0]
                    basis1 = cont1[1]
                    ss0 = ['-','+','+','(',')',',']
                    ss1 = ['\-','\+','\+','(',')',',']
                    ss2 = ['','j','j','','','']
                    for js,sj in enumerate(ss1):
                        if ss0[js] in basis1:
                            cmd = "echo '%s' | sed -n 's/%s/%s/p'"%(basis1, sj, ss2[js])
                            basis1 = os.popen(cmd).read().strip()
                    break

        if method1 is None:
            print ' #ERROR: method cannot be retrieved??'
            sys.exit(3)
        #else:
        #    if 'maxcyc' in method1:
        #        # e.g.,  "# uqcisd(t)(maxcyc=100)/6-311++g(d,p)"
        #        idx0 = method1.index('(max')
        #        method1 = method1[:idx0]

        # now read energy
        m1U = method1.upper(); #print ' ** meth = ', m1U
        if m1U in ['B3LYP', 'RB3LYP', 'UB3LYP']:
            cmd2 = "grep 'E([RU]B3LYP) =' %s | tail -n 1 | awk '{print $5}'"%f #out
            # e.g.,
            #  SCF Done:  E(RB3LYP) =  -40.5236797399     A.U. after   10 cycles
            #print cmd2
            E_str = os.popen(cmd2).read().strip()
        elif m1U in ['CCSD(T)', 'QCISD(T)', 'UCCSD(T)', 'UQCISD(T)']:
            m1Uu = {'UQCISD(T)':'QCISD(T)', 'QCISD(T)':'QCISD(T)', 'UCCSD(T)':'CCSD(T)', 'CCSD(T)':'CCSD(T)'}[m1U]
            cmd2 = "grep '^ %s=' %s"%(m1Uu,f); #print cmd2
            # e.g.,
            #  QCISD(T)= -0.11358793642D+03
            E_str = os.popen(cmd2).read().strip().split('\n')[0].split()[-1]
        elif m1U in ['HF','RHF','UHF',]:
            cmd2 = "grep 'E([RU]HF) =' %s | tail -n 1"%f
            E_str = os.popen(cmd2).read().strip().split('=')[1].split()[0]
        elif m1U in ['MP2','RMP2','UMP2',]:
            cmd2 = "grep 'EUMP2 = ' %s | tail -n 1"%f
            E_str = os.popen(cmd2).read().strip().split()[-1]
        else:
            print '#ERROR: not implemented yet'; sys.exit(3)

        if 'D' in E_str:
            E = eval( 'E'.join( E_str.split('D') ) )
        else:
            E = eval( E_str )

        self.E = E*self.const

        atoms = self.atoms
        qcl = method1 + basis1
        self.qcl = qcl
        #print '  -- qcl = ', qcl
        if etype == 'AE':
            ea_refs = data.retrieve_esref( qcl.lower() )
            E0 = 0.
            symbs = [ ai.symbol for ai in atoms ]
            for si in symbs: E0 += ea_refs[si][0]
            self.AE = (E - E0)*self.const
            self.update_dic( ['qcl','AE'], [self.qcl, self.AE] )
        elif etype == 'EE': # electronic energy
            cmd1 = "awk '/nuclear repulsion energy/{print NR}' %s | tail -n 1"%f
            En = eval( io2.cmdout2(cmd1) ) * self.const
            Ee = E - En
            self.En = En
            self.Ee = Ee
            self.update_dic( ['qcl','E','EN','EE'], [self.qcl,self.E,En,Ee] )
        else:
            self.update_dic( ['qcl','E',], [self.qcl, self.E] )

    def get_mo_energy(self):
        """
        get HOMO, LUMO and their gap
        """
        f = self.f
        cmd1 = "awk '/^ The electronic state is/{print NR}' %s | tail -n 1"%f
        Ln1 = int(io2.cmdout2(cmd1)) + 1
        cmd2 = "awk '/^          Condensed to atoms \(all electrons\)/{print NR}' %s | tail -n 1"%f
        Ln2 = int(io2.cmdout2(cmd2)) - 1
        cmd = "sed -n '%d,%dp' %s | grep Beta"%(Ln1,Ln2,f)
        cont0 = io2.cmdout2(cmd)
        if cont0 != '':
            print ' Alpha & Beta spins are both involved, spin polarized??'
            sys.exit(2)
        else:
            cmd = "sed -n '%d,%dp' %s | grep ' Alpha virt. eigenvalues --' | head -n 1"%(Ln1,Ln2,f)
            ct = io2.cmdout2(cmd);
            self.lumo = eval( ct.split()[-5] ) * self.uo.h2e
            cmd = "sed -n '%d,%dp' %s | grep ' Alpha  occ. eigenvalues --' | tail -1"%(Ln1,Ln2,f)
            ct = io2.cmdout2(cmd)
            self.homo = eval( ct.split()[-1] ) * self.uo.h2e
            #iseed = debug(iseed)
            self.gap = self.lumo - self.homo

            self.update_dic(['homo','lumo','gap'], [self.homo, self.lumo, self.gap])
            self.update_dic(['HOMO','LUMO','GAP'], [self.homo, self.lumo, self.gap])

    def get_thermo(self, scale_factor=0.965):
        """
        for DFT, the `scale_factor is 0.965
        """
        f = self.f
        cmd = "grep '^ Freq' %s"%f
        # data in lines below could be retrieved together
        #
        # Sum of electronic and zero-point Energies=           -174.095490
        # Sum of electronic and thermal Energies=              -174.091365
        # Sum of electronic and thermal Enthalpies=            -174.090421
        # Sum of electronic and thermal Free Energies=         -174.122620
        conts = io2.cmdout(cmd)
        assert conts != [], '#ERROR: no thermochem was done!'
        if not isGn:
            cmd = "awk '/^ Sum of electronic and/{print $NF}' %s"%f
            U0, U, H, G = np.array( io2.cmdout(cmd) ).astype(np.float) * self.const
            self.U0, self.U, self.H, self.G  = U0, U, H, G

            cmd = "awk '/^ Zero-point correction=/{print $3}' %s"%f
            zpe = eval(io.cmdout(cmd)[0]) * self.const # the last two entries are "0.81475823 Hartree/atom"
            self.zpe = self.zpve = zpe
        else:
            print '#ERROR: cannot handle output of Compositional method like G4MP2 yet!'
            sys.exit(2)
        E_ = U0 - zpe
        #assert abs(E_ - E) < 0.0001, '#ERROR: '

        # people usually scale ZPE by a factor, for B3LYP, it's 0.965??
        U0c = E_ + scale_factor*zpe
        self.U0c = U0c
        self.update_dic(['U0', 'U0c', 'U', 'H', 'G'], [U0, U0c, U, H, G])

    def get_dipole_moment(self):
        cmd = "grep -A1 ' Dipole moment (field-independent basis, Debye):' %s | tail -1 | awk '{print $NF}'"%self.f
        self.mu = eval( io2.cmdout2(cmd) )
        self.update_dic(['MU', 'DIPOLE'], [self.mu, self.mu])

    def get_polarizability(self):
        cmd = "awk '/Isotropic polarizability/{print $6}' %s"%self.f
        #print cmd
        try:
            self.alpha = eval( io2.cmdout2(cmd) )
            self.update_dic(['ALPHA',], [self.alpha,])
        except:
            print ' * WARNING: no Isotropic polarizability found'

    def get_nmr(self):
        cmd = "awk '/  Isotropic =  /{print $5}' %s"%self.f
        vals = io2.cmdout(cmd)
        self.nmr = [ eval(val) for val in vals ]
        self.update_dic(['NMR',], [self.nmr,])

    def get_mulliken_population(self):
        # note that the output Mulliken charge starting line may be different
        # for different versions of Gaussian, i.e., with or without `atomic` in between
        # "Mullike" and "charges:", so a regular expression is used
        cmd1 = "awk '/^ Mulliken [a-zA-Z]*\s?charges:/{print NR}' %s | tail -1"%self.f
        #print cmd1
        Ln1 = int(io2.cmdout2(cmd1)) + 2
        cmd2 = "awk '/^ Sum of Mulliken [a-zA-Z]*\s?charges =/{print NR}' %s | tail -1"%self.f
        Ln2 = int(io2.cmdout2(cmd2)) - 1
        cmd = "sed -n '%d,%dp' %s"%(Ln1,Ln2,self.f)
        cs = io2.cmdout2(cmd).split('\n')
        pops = [ eval(ci.split()[2]) for ci in cs ]
        #pops = np.array(vals)
        self.populations = pops
        self.update_dic(['POPULATION',], [self.populations,])

    def get_force(self):
        iou = io2.Units()
        const = iou.h2e / iou.b2a # from hartree/bohr to eV/A
        #cmd = "grep '^\s*[XYZ][0-9]*   ' %s | awk '{print $3}'"%self.f #
        cmd1 = "awk '/^ Variable       Old X    -DE/{print NR}' %s | tail -1"%self.f
        Ln1 = int(io2.cmdout2(cmd1)) + 2
        cmd2 = "awk '/^         Item               Value     Threshold  Converged/{print NR}' %s | tail -1"%self.f
        Ln2 = int(io2.cmdout2(cmd2)) - 1
        cmd = "sed -n '%d,%dp' %s"%(Ln1,Ln2,self.f)
        #print cmd
        cs = io2.cmdout2(cmd).split('\n')
        vals = [ eval(ci.split()[2]) for ci in cs ]
        #print vals
        #print len(vals)
        forces = np.array(vals).reshape((self.na, 3)) * const
        #abs_forces = np.linalg.norm( forces, axis=1 )
        #self.forces = abs_forces[:, np.newaxis]
        self.forces = forces
        self.update_dic(['FORCE',], [self.forces,])

class GRs(object):

    def __init__(self, fs, properties=['AE'], unit='kcal', write_Y=False, nproc=1, istart=0):

        self.n = len(fs)
        typeP = type(properties)
        if typeP is str:
            properties = [ properties.upper(), ]
        elif typeP is list:
            properties = [ prop.upper() for prop in properties ]
        else:
            raise '#ERROR,'
        npr = len(properties)
        istats = [] # is_local_property = False
        for propi in properties:
            istat = False # assume global property
            if propi in ['POPULATION','NMR']:
                istat = True
            istats.append( istat )

        if nproc == 1:
            self.objs = []
            for i,f in enumerate(fs):
                #print i+1, f
                ipt = [ f, properties, unit, istart ]
                self.objs.append( self.processInput(ipt) )
        else:
            pool = multiprocessing.Pool(processes=nproc)
            ipts = [ [ fi, properties, unit ] for fi in fs ]
            self.objs = pool.map(self.processInput, ipts)

        ys = []
        #print ' - npr = ', npr
        for ipr in range(npr):
            ysi = []
            for obj in self.objs:
                #print ' - obj.properties = ', obj.properties
                yi = obj.properties[ipr]; #print ' ysi = ', ysi
                if istats[ipr]: # atomic property, e.g., NMR shifts
                    ysi += yi
                else:
                    ysi.append( yi )
            ys.append( ysi )
        #print ys
        self.dic = dict( zip(properties, ys) )
        if write_Y:
            for ipr in range(npr):
                np.savetxt('%s.dat'%properties[ipr], ys[ipr], fmt='%.2f')

    def processInput(self, ipt):
        f, properties, unit, istart = ipt
        obj = GR(f, properties=properties, unit=unit, istart=istart)
        return obj

    def get_statistics(self):
        """
        `zs, `nas, `nhass, `coords, etc
        """
        zs = []
        zsr = []
        nas = []
        nhass = []
        zsu = set([])
        coords = []
        for i in range(self.n):
            obj_i = self.objs[i]
            coords.append( obj_i.coords )
            zsi = obj_i.zs
            nhass.append( (np.array(zsi)>1).sum() )
            nas.append( len(zsi) )
            zsu.update( zsi )
            zsr += list(zsi)
            zs.append( zsi )
        zsu = list(zsu)
        nzu = len(zsu)
        zsu.sort()
        nzs = np.zeros((self.n, nzu), np.int32)
        for i in range(self.n):
            for iz in range(nzu):
                istats = np.array(zs[i]) == zsu[iz]
                nzs[i,iz] = np.sum(istats)
        self.nzs = nzs
        self.zsu = zsu
        self.zs = zs
        self.zsr = np.array(zsr,np.int32)
        self.nas = np.array(nas,np.int32)
        self.nhass = np.array(nhass,np.int32)
        self.coords = coords

def get_line_number(cmd):
    return int(io2.cmdout2(cmd).split(':')[0])

def read_coords(f):
    lines = file(f).readlines()
    cmd0 = "grep -n ' Input orientation:' %s | tail -1"%f
    #cmd0 = "grep -n ' Standard orientation:' %s | tail -1"%f
    #print cmd0
    ln = get_line_number(cmd0) + 5
    zs = []; coords = []
    while 1:
        l = lines[ln-1].strip()
        if l[:10] == '-'*10: break
        ia,zi,_,x,y,z = l.split()
        zs.append(int(zi))
        coords.append([eval(xi) for xi in [x,y,z]])
        ln += 1
    return zs,coords

    obsolete="""#cmd1 = "grep -n '    Distance matrix (angstroms):' %s | tail -1"%f
    cmd1 = "grep -n ' Rotational constants (GHZ):' %s"%f
    #print cmd1
    ln1s = get_line_number(cmd1);
    if type(ln1s) is int: ln1s = [ln1s,]
    nnl = len(ln1s)
    fconts = file(f).readlines()
    ln1_found = False
    for jn in range(1,nnl+1):
        if fconts[ ln1s[-jn]-2 ].strip() == '-'*69:
            ln1 = ln1s[-jn] - 2
            ln1_found = True
            break

    if not ln1_found:
        print ' #ERROR: `ln1 not assigned!!'"""

def get_cf_and_chg(f):
    """ cf: chemical formula
        chg: charge """
    cmdi = "grep Stoichiometry %s | tail -1"%f
    ct = io2.cmdout2(cmdi)
    info = ct.split()[-1]
    if ('(' in info) and (info[-1] == ')'):
        # charged mol
        cf, c0_ = info.split('(') # cf: chemical formula;
        if c0_[-2] == '-':
            chg = -int(c0_[:-2]) # negatively charged
        elif c0_[-2] == '+':
            chg = int(c0_[:-2]) # positively charged
        else:
            #print 'c0_ = ', c0_, ', ct = ', ct
            #sys.exit(2)
            chg = 0 # spin-polarized case, e.g., C4H9O(2)
    else:
        chg = 0
        cf = info
    return cf, chg

def get_spin(f):
    cmdi = "grep 'Multiplicity =' %s | tail -1"%f
    ct = cmdout2(cmdi).split()
    chg = int(ct[2]); sm = int(ct[-1])
    return sm, chg

if __name__ == "__main__":
    import sys
    import stropr as so

    args = sys.argv[1:]

    idx = 0
    keys=['-p','-properties']; hask,s,idx = so.parser(args,keys,'E',idx)
    assert hask
    props = s.split(',')

    fs = args[idx:]

    #test
    obj = GRs(fs, properties=props)
    for propi in props:
        for (f,p) in zip(fs,obj.dic[propi]):
            print f,p