#!/usr/bin/python3

"""Swing molecules based on normal coordinates."""

import os
import argparse
import shutil
from collections import defaultdict

import numpy as np
import matplotlib.pyplot as plt
from ase import Atoms
from ase.calculators.orca import ORCA
from cclib import ccopen
from scipy.constants import physical_constants
from scipy.optimize import root
from scipy.optimize import minimize
from scipy.optimize import minimize_scalar

try:
    from xtb.ase.calculator import XTB
except ModuleNotFoundError as e:
    print("WARNING: won't be able to use XTB")

hartree, _, _ = physical_constants["Hartree energy"]

smd2gbsa = defaultdict(
    lambda: "none",
    {
        "acetone": "acetone",
        "mecn": "acetonitrile",
        "acetonitrile": "acetonitrile",
        "benzene": "benzene",
        "dichloromethane": "ch2cl2",
        "chloroform": "chcl3",
        "carbon disulfide": "cs2",
        "dmf": "dmf",
        "n,n-dimethylformamide": "dmf",
        "dmso": "dmso",
        "dimethylsulfoxide": "dmso",
        "diethyl ether": "ether",
        "water": "water",
        "methanol": "methanol",
        "n-hexane": "n-hexan",
        "thf": "thf",
        "tetrahydrofuran": "thf",
        "toluene": "toluene",
    },
)

def main():
    """Run main procedure."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("logfile")
    # TODO(schneiderfelipe): set charge and multiplicity
    parser.add_argument(
        "-a",
        "--acc",
        help="accuracy for SCC calculation, lower is better",
        type=float,
        default=1.0,
    )
    parser.add_argument(
        "--iterations",
        help="number of iterations in SCC",
        type=int,
        default=250,
    )
    parser.add_argument(
        "--gfn", help="specify parametrisation of GFN-xTB", type=int
    )
    parser.add_argument(
        "--etemp", help="electronic temperature", type=float, default=300.0
    )
    parser.add_argument(
        "-s",
        "--solvent",
        help=("solvent (SMD/GBSA implicit solvation models)"),
        default="none",
    )
    parser.add_argument(
        "--do-not-cache-api",
        dest="cache_api",
        help="Do not reuse generate API objects (not recommended)",
        action="store_false",
    )

    parser.add_argument(
        "--pm3", help="use PM3", action="store_true",
    )
    parser.add_argument(
        "--b97-3c", help="use B97-3c", action="store_true",
    )
    parser.add_argument(
        "--minimize", action="store_true",
    )
    parser.add_argument(
        "--transition-state", action="store_true",
    )
    parser.add_argument("--max-omega", type=float, default=1.0)
    parser.add_argument("--tol", type=float, default=1e-3)
    parser.add_argument("--nprocs", type=int, default=4)
    args = parser.parse_args()
    print(args)

    data = ccopen(args.logfile).parse()
    initial_positions = data.atomcoords[-1]
    atoms = Atoms(numbers=data.atomnos, positions=initial_positions)

    if args.gfn:
        method = f"GFN{args.gfn}-xTB"
        solvent = smd2gbsa[args.solvent.lower()]

        calc = XTB(
            method=method,
            accuracy=args.acc,
            electronic_temperature=args.etemp,
            max_iterations=args.iterations,
            solvent=solvent,
            cache_api=args.cache_api,
        )
    else:

        if args.b97_3c:
            method = "B97-3c D3BJ def2-SV(P)"
        elif args.pm3:
            method = "PM3"
        else:

            def allow_keyword(keyword):
                for forbidden in {"freq", "opt", "irc", "print"}:
                    if forbidden in keyword.lower():
                        return False
                return True

            keywords = [
                keyword
                for keyword in data.metadata["keywords"]
                if allow_keyword(keyword)
            ]

            method = " ".join(keywords)

        solvent = args.solvent
        blocks = f"%pal\n nprocs {args.nprocs}\nend\n%scf\n maxiter {args.iterations}\nend"
        if solvent != "none" and not args.pm3:
            blocks += f'\n%cpcm\n smd true\n smdsolvent "{solvent}"\nend'

        if "ORCA_COMMAND" not in os.environ:
            # For parallel runs ORCA has to be called with full pathname
            os.environ["ORCA_COMMAND"] = shutil.which("orca")

        calc = ORCA(
            label="012345_swing", orcasimpleinput=method, orcablocks=blocks
        )

    print(f"*** {method} ***")
    print(f"    : solvent:              {solvent}")

    atoms.set_calculator(calc)
    potential_min = atoms.get_potential_energy()
    print(f"@ potential energy:        {potential_min} eV")

    indices = np.where(data.vibfreqs < 0)[0]
    n_indices = len(indices)
    print(f"@ imaginary frequencies:   {data.vibfreqs[indices]}")
    if not n_indices:
        print("    : nothing to be done, bye")
        return

    ignoring = None
    if args.transition_state:
        ignoring = 0
        print("    : transition state:    ignoring first imaginary frequency")

    omegas = []
    potentials = []

    def f(omega):
        atoms.set_positions(
            initial_positions
            + np.einsum("i,ijk->jk", omega, data.vibdisps[indices])
        )

        potential = 1e3 * (atoms.get_potential_energy() - potential_min)

        omegas.append(omega)
        potentials.append(potential)
        print(f"    : omega:               {omega}")
        print(f"    : potential:           {potential} meV")

        return potential

    if args.minimize:
        guesses = [np.zeros_like(indices, dtype=float)]

    for i in indices:
        if ignoring is not None and i == ignoring:
            continue

        print(f"@ searching in direction   #{i}")

        def g(w):
            z = np.zeros_like(indices, dtype=float)
            z[i] = w
            return f(z)

        if args.minimize:
            res = minimize_scalar(
                g,
                method="bounded",
                bounds=(-args.max_omega, args.max_omega),
                tol=args.tol,
            )
            print(res)

            guess = np.zeros_like(indices, dtype=float)
            guess[i] = res.x
            guesses.append(guess)
        else:
            dx = args.max_omega / 100
            x = [-dx, 0.0, dx]
            y = [g(-dx), 0.0, g(dx)]

            # p[0] * x**2 + p[1] * x + p[2] == k * (x - x0)**2 == k * x**2 - 2 * x0 * k * x + k * x0**2
            p = np.polyfit(x, y, 2)
            print(p)
            print(np.roots(p))

            dp = np.polyder(p)
            print(dp)
            r = np.roots(dp)
            print(r)

            # k = p[0]
            # x0 = np.sqrt(p[2] / k)
            # print(k, x0)
            # print(root(lambda z: [p[0] - z[0], p[1] + 2 * z[0] * z[1], p[2] - z[0] * z[1] ** 2], [k, x0]))

            best_positions = initial_positions + np.einsum(
                "i,ijk->jk", r, data.vibdisps[indices]
            )

    if args.minimize:
        print("@ choosing initial guess for global search")
        if n_indices > 1:
            guesses.append(np.sum(guesses, axis=0))
        x0 = guesses[np.argmin([f(guess) for guess in guesses])]

        print("@ searching in all directions")
        constraints = ()
        if args.transition_state and ignoring is not None:
            constraints = (
                {"type": "eq", "fun": lambda omega: omega[ignoring]},
            )
        res = minimize(
            f,
            x0=x0,
            bounds=n_indices * [(-args.max_omega, args.max_omega)],
            constraints=constraints,
            tol=args.tol,
        )
        print(res)
        best_positions = initial_positions + np.einsum(
            "i,ijk->jk", res.x, data.vibdisps[indices]
        )

        # TODO(schneiderfelipe): correct for when using --transition-state
        omegas = np.array(omegas)
        fig, ax = plt.subplots(n_indices, 1)
        if n_indices == 1:
            ax = [ax]
        xlim = (-args.max_omega - 0.05, args.max_omega + 0.05)
        ylim = (np.min(potentials) - 2.0, 40.0)
        for i in indices:
            if ignoring is not None and i == ignoring:
                continue

            ax[i].plot(omegas[:, i], potentials, "o")
            ax[i].set_title(f"view of normal mode #{i}")
            ax[i].set_ylabel(r"potential energy (meV)")
            ax[i].set_xlabel(rf"$\omega_{i}$")
            ax[i].set_ylim(ylim)
            ax[i].set_xlim(xlim)
        plt.tight_layout()
        plt.show()

    print("@ writing best geometry to swinged.xyz")
    # TODO(schneiderfelipe): print a RMSD between initial and final structures
    atoms.set_positions(best_positions)
    atoms.write("swinged.xyz", format="xyz", plain=True)

if __name__ == "__main__":
    main()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jul 20 9:35 2020

@author: MCR

File containing the necessary functions to create an empirical
interpolated trace model in the overlap region for SOSS order 1.
"""

import os
import numpy as np
from astropy.io import fits
from scipy.optimize import least_squares
import webbpsf
from SOSS.trace import tracepol as tp
from SOSS.extract.empirical_trace import centroid as ctd
from SOSS.extract.empirical_trace import plotting as plotting
from SOSS.dms.soss_engine import ExtractionEngine
from SOSS.extract.engine_legacy import ThroughputSOSS, WebbKer

def build_empirical_trace(clear, F277W, filename='spatial_profile.fits',
                          pad=(0, 0), doplot=False, verbose=False):
    ''' Procedural function to wrap around construct orders 1 and 2.
    Will do centroiding and call the functions to construct the models.

    ***Will eventually want clear and F277W to be the full fits with headers
    to get parameters from***

    Parameters
    ----------
    clear : np.array of float (2D) - eventually path
    F277W : np.array of float (2D) - eventually path
    filename : str
    pad : tuple
    doplot : bool
    verbose : bool
    '''

    # Print overwrite warning if output file already exists.
    if os.path.exists(filename):
        print('Output file {} already exists. It will be overwritten'.format(filename))
    # Get the centroid positions for both orders from the data.
    if verbose is True:
        print('Getting trace centroids...')
    centroids, rot_pars = ctd.get_contam_centroids(clear, doplot=doplot,
                                                   bound=False,
                                                   showprogress=verbose)

    # Overplot the data centroids on the CLEAR exposure if desired.
    if doplot is True:
        plotting._plot_centroid(clear, centroids)

    # Construct the first order profile.
    if verbose is True:
        print('Interpolating the first order trace...')
    o1frame = construct_order1(clear, F277W, rot_pars, centroids, pad=pad[0],
                               doplot=doplot)
    # Pad the spectral axis.
    if pad[1] != 0:
        if verbose is True:
            print('Adding padding to first order spectral axis...')
        o1frame = pad_spectral_axis(o1frame, centroids['order 1'][0],
                                    centroids['order 1'][1], pad=pad[1])

    # Get the extraction parameters
    #extract_params = get_extract_params()
    #ref_file_args = get_ref_file_args(o1frame)
    # Construct the second order profile.
    #o2frame_contam = construct_order2(clear, ref_file_args, extract_params)
    # Create a mask to remove residuals from the first order.
    #o2frame = mask_order(o2frame_contam, x2, y2)
    # Set any spurious negative values to zero.
    #o2frame[o2frame < 0] = 0
    # Normalize the profile in each column.
    #o2frame /= np.nansum(o2frame, axis=0)

    # Write the trace model to disk.
    #hdu = fits.PrimaryHDU()
    #hdu.data = np.dstack((o1frame, o2frame))
    #hdu.writeto(filename, overwrite=True)

    if verbose is True:
        print('Done.')

    return o1frame#, o2frame

def calc_interp_coefs(make_psfs=False, doplot=True, F277W=True, filepath=''):
    '''Function to calculate the interpolation coefficients necessary to
    construct a monochromatic PSF profile at any wavelength between
    the two 1D PSF anchor profiles. Linear combinations of the blue and red
    anchor are iteratively fit to each intermediate wavelength to find the
    best fitting combination. The mean linear coefficients across the 10 WFE
    error realizations are returned for each wavelengths.
    When called, 2D moncochromatic PSF profiles will be generated and saved
    to disk if the user does not already have them available.
    This should not need to be called by the end user except in rare cases.

    Parameters
    ----------
    make_psfs : bool
        Whether or not WebbPSF will have to generate the monochromatic
        PSFs used for the fitting.
    doplot : bool
        Whether to show the diagnostic plots for the model derivation.
    F277W : bool
        Set to False if no F277W exposure is available for the observation.
        Finds coefficients for the entire 2.1 - 2.9µm region in this case.
    filepath : str
        Path to directory containing the WebbPSF monochromatic PSF fits
        files, or the directory to which they will be stored when made.
        Defaults to the current directory.

    Returns
    -------
    pb : np.array of float
        Polynomial coefficients of the interpolation index fits for the
        blue anchor.
    pr : np.array of float
        Polynomial coefficients of the interpolation index fits for the
        red anchor.
    '''

    # Red anchor is 2.8µm without an F277W exposure.
    if F277W is False:
        wave_range = np.linspace(2.1, 2.9, 7)
    # Red anchor is 2.5µm with F277W exposure.
    else:
        wave_range = np.linspace(2.1, 2.45, 7)

    # Read in monochromatic PSFs generated by WebbPSF.
    PSFs = []
    # Loop over all 10 available WFE realizations.
    for i in range(10):
        psf_run = []
        # Create the PSF if user has indicated to.
        if make_psfs is True:
            loicpsf(wavelist=wave_range*1e-6, wfe_real=i)
        # If user already has PSFs generated.
        for w in wave_range:
            try:
                psf_run.append(fits.open('{0:s}SOSS_os10_128x128_{1:.6f}_{2:.0f}.fits'
                                         .format(filepath, w, i))[0].data)
            # Generate missing PSFs if necessary.
            except FileNotFoundError:
                print('No monochromatic PSF found for {0:.2f}µm and WFE realization {1:.0f}.'
                      .format(w, i))
                loicpsf(wavelist=[w*1e-6], wfe_real=i, filepath=filepath)
                psf_run.append(fits.open('{0:s}SOSS_os10_128x128_{1:.6f}_{2:.0f}.fits'
                                         .format(filepath, w, i))[0].data)
        PSFs.append(psf_run)

    # Determine specific interpolation coefficients for all WFEs
    wb, wr = [], []

    for E in range(10):
        # Generate the blue wavelength anchor.
        # The width of the 1D PSF has lambda/D dependence, so rescale all
        # profiles to a common wavelength to remove these chromatic effects.
        rngeb = np.linspace(0, round(1280*(2.5/2.1), 0) - 1, 1280)
        offsetb = rngeb[640] - 640
        newb = np.interp(np.arange(1280), rngeb - offsetb,
                         np.sum(PSFs[E][0][600:700, :], axis=0))

        # Generate the red wavelength anchor.
        if F277W is False:
            # Use 2.85µm for CLEAR.
            rnger = np.linspace(0, round(1280*(2.5/2.9), 0) - 1, 1280)
        else:
            # Or 2.42µm for F277W.
            rnger = np.linspace(0, round(1280*(2.5/2.45), 0) - 1, 1280)
        offsetr = rnger[640] - 640
        # Remove lambda/D scaling.
        newr = np.interp(np.arange(1280), rnger - offsetr,
                         np.sum(PSFs[E][6][600:700, :], axis=0))

        # Loop over all monochrmatic PSFs to determine interpolation coefs.
        for f, wave in enumerate(wave_range):
            # Lists for running counts of indicies and model residuals.
            resid, ind_i, ind_j = [], [], []

            # Rescale monochrmotic PSF to remove lambda/D.
            newrnge = np.linspace(0, round(1280*(2.5/wave), 0) - 1, 1280)
            newoffset = newrnge[640] - 640
            newpsf = np.interp(np.arange(1280), newrnge - newoffset,
                               np.sum(PSFs[E][f][600:700, :], axis=0))

            # Brute force the coefficient determination.
            for i in range(1, 100):
                for j in range(1, 100):
                    i /= 10
                    j /= 10
                    # Create a test profile which is a mix of the two anchors.
                    mix = (i*newb + j*newr) / (i + j)

                    # Save current iteration in running counts.
                    resid.append(np.sum(np.abs(newpsf - mix)[450:820]))
                    ind_i.append(i)
                    ind_j.append(j)

            # Determine which combination of indices minimizes model residuals.
            ind = np.where(resid == np.min(resid))[0][0]
            # Save the normalized indices for this wavelength.
            wb.append(ind_i[ind] / (ind_i[ind] + ind_j[ind]))
            wr.append(ind_j[ind] / (ind_i[ind] + ind_j[ind]))

    wb = np.reshape(wb, (10, 7))
    wr = np.reshape(wr, (10, 7))

    # Fit a second order polynomial to the mean of the interpolation indices.
    pb = np.polyfit(wave_range, np.mean(wb, axis=0), 2)
    pr = np.polyfit(wave_range, np.mean(wr, axis=0), 2)

    # Show the diagnostic plot if necessary.
    if doplot is True:
        plotting._plot_interpmodel(wave_range, wb, wr, pb, pr)

    return pb, pr, wb, wr

def _chromescale(wave_start, profile, ycen, wave_end=2.5, invert=False):
    '''Remove the lambda/D chromatic PSF scaling by re-interpolating a
    monochromatic PSF function onto a rescaled axis.

    Parameters
    ----------
    wave_start : float
        Wavelength corresponding to the input 1D PSF profile.
    profile : np.array of float
        1D PSF profile to be rescaled.
    ycen : float
        Y-pixel position of the order 1 trace centroid.
    wave_end : float
        Wavelength after rescaling.
    invert : bool
        If True, add back the lambda/D scaling instead of removing it.

    Returns
    -------
    prof_recsale : np.array of float
        Rescaled 1D PSF profile.
    '''

    xrange = len(profile)
    # Assume that the axis corresponding to wave_end can be constructed with
    # np.arange - this is the 'standard' axis.
    # Construct the input profile axis by rescaling the standard axis by the
    # ratio of wave_end/wave_start.
    xax = np.linspace(0, round(xrange*(wave_end/wave_start), 0) - 1, xrange)
    # Calculate the offset of the centroid y-position in the standard and
    # rescaled axis, to ensure it remains at the same pixel position.
    offset = xax[int(round(ycen, 0))] - ycen

    # Interpolate the profile onto the standard axis.
    if invert is False:
        prof_rescale = np.interp(np.arange(xrange), xax - offset, profile)
    # Or interpolate the profile from the standard axis to re-add
    # the lambda/D scaling.
    else:
        prof_rescale = np.interp(xax - offset, np.arange(xrange), profile)

    return prof_rescale

def construct_order1(clear, F277, rot_params, ycens, pad=0, doplot=False):
    '''This creates the full order 1 trace profile model. The region
    contaminated by the second order is interpolated from the CLEAR and F277W
    exposures, or just from the CLEAR exposure and a standard red anchor if
    no F277W exposure is available.
    The steps are as follows:
        1. Determine the red and blue anchor profiles for the interpolation.
        2. Construct the interpolated profile at each wavelength in the
           overlap region.
        3. Stitch together the original CLEAR exposure and the interpolated
           trace model (as well as the F277W exposure if provided).
        4. Mask the contamination from the second and third orders, and
           reconstruct the underlying wing structure of the first order -
           including any padding in the spatial direction.

    Parameters
    ----------
    clear : np.array of float (2D)
        NIRISS SOSS CLEAR exposure dataframe.
    F277 : np.array of float (2D)
        NIRISS SOSS F277W filter exposure dataframe. If no F277W exposure
        is available, pass None.
    rot_params : list of float
        List containing the rotation angle, X and Y anchor points, and X
        and Y offset required to transform OM coordinates to the detector
        frame.
    ycens : dict
        Dictionary of Y-coordinates for the trace centroids of the first three
        diffraction orders, ie. as returned by get_contam_centroids.
    pad : int
        Number of pixels of padding to add on both ends of the spatial axis.
    doplot : bool
        if True, do diagnostic plotting.

    Returns
    -------
    newmap : np.array of float (2D)
        Interpolated order 1 trace model with padding.
    '''

    # Determine the anchor profiles - blue anchor.
    # Note: Y-values returned from OM are inverted relative to the UdeM
    # coordinate system.
    ###### TODO Switch this to use wavecal ######
    xOM, yOM, tp2 = ctd.get_om_centroids()
    # Get OM centroid pixel coords at 2.1µm.
    xom22 = tp.wavelength_to_pix(2.1, tp2, 1)[0]
    yom22 = 256 - tp.wavelength_to_pix(2.1, tp2, 1)[1]
    # Use rot params to find location of the 2.1µm centroids in the data frame.
    xd22, yd22 = ctd.rot_centroids(*rot_params, xom22, yom22)
    xd22 = np.round(xd22, 0).astype(int)[0]
    yd22 = np.round(yd22, 0).astype(int)[0]
    # Extract the 2.1µm anchor profile from the data.
    Banch = clear[:, xd22]
    # Mask second and third order, reconstruct wing structure and pad.
    cens = [ycens['order 1'][1][xd22], ycens['order 2'][1][xd22],
            ycens['order 3'][1][xd22]]
    Banch = reconstruct_wings(Banch, ycens=cens, contamination=True, pad=pad,
                              doplot=doplot)
    # Remove the lambda/D scaling.
    Banch = _chromescale(2.1, Banch, yd22)
    # Normalize
    Banch /= np.nansum(Banch)

    # Determine the anchor profiles - red anchor.
    if F277 is None:
        # If no F277W exposure is provided, interpolate out to 2.9µm.
        # Generate a simulated 2.9µm PSF.
        stand = loicpsf([2.9*1e-6], save_to_disk=False, oversampling=1,
                        pixel=256, verbose=False)[0][0].data
        # Extract the spatial profile.
        Ranch = np.sum(stand[124:132, :], axis=0)
        # Detemine OM coords of 2.9µm centroid.
        xom25 = tp.wavelength_to_pix(2.9, tp2, 1)[0]
        yom25 = 256 - tp.wavelength_to_pix(2.9, tp2, 1)[1]
        # Transform into the data frame.
        yd25 = ctd.rot_centroids(*rot_params, xom25, yom25, bound=False)[1]
        yd25 = int(round(yd25[0], 0))
        # Interpolate the WebbPSF generated profile to the correct location.
        Ranch = np.interp(np.arange(256), np.arange(256)-128+yd25, Ranch)
        # Reconstruct wing structure and pad.
        Ranch = reconstruct_wings(Ranch, ycens=[yd25], contamination=False,
                                  pad=pad)
        # Rescale to remove chromatic effects.
        Ranch = _chromescale(2.9, Ranch, yd25)
        # Normalize
        Ranch /= np.nansum(Ranch)

        # Interpolation polynomial coeffs, calculated via calc_interp_coefs
        coef_b = [0.80175603, -5.27434345, 8.54474316]
        coef_r = [-0.80175603, 5.27434345, -7.54474316]
        # Pixel coords at which to start and end interpolation in OM frame.
        end = int(round(tp.wavelength_to_pix(2.1, tp2, 1)[0], 0))
        # Determine the OM coords of first pixel on detector
        start = int(round(xom25, 0))
        rlen = end - start

    else:
        # If an F277W exposure is provided, only interpolate out to 2.45µm.
        # Redwards of 2.45µm we have perfect knowledge of the order 1 trace.
        # Get OM centroid pixel coords at 2.45µm
        ###### TODO Switch this to use wavecal ######
        xom25 = tp.wavelength_to_pix(2.45, tp2, 1)[0]
        yom25 = 256 - tp.wavelength_to_pix(2.45, tp2, 1)[1]
        # Transform into the data frame.
        xd25, yd25 = ctd.rot_centroids(*rot_params, xom25, yom25)
        xd25 = np.round(xd25, 0).astype(int)[0]
        yd25 = np.round(yd25, 0).astype(int)[0]
        # Extract and rescale the 2.5µm profile.
        Ranch = F277[:, xd25-1]
        # Reconstruct wing structure and pad.
        cens = [ycens['order 1'][1][xd25-1]]
        Ranch = reconstruct_wings(Ranch, ycens=cens, contamination=False,
                                  pad=pad)
        Ranch = _chromescale(2.45, Ranch, yd25)
        # Normalize
        Ranch /= np.nansum(Ranch)

        # Interpolation polynomial coeffs, calculated via calc_interp_coefs
        coef_b = [1.51850915, -9.76581613, 14.80720191]
        coef_r = [-1.51850915,  9.76581613, -13.80720191]
        # Pixel coords at which to start and end interpolation in OM frame.
        end = int(round(tp.wavelength_to_pix(2.1, tp2, 1)[0], 0))
        start = int(round(tp.wavelength_to_pix(2.45, tp2, 1)[0], 0))
        rlen = end - start

    # Create the interpolated order 1 PSF.
    map2D = np.zeros((256+2*pad, 2048))*np.nan
    # Get OM X-pixel values for the region to be interpolated.
    cenx_om = np.arange(rlen) + start
    # Find the wavelength at each X centroid
    lmbda = tp.specpix_to_wavelength(cenx_om, tp2, 1)[0]
    # Y centroid at each wavelength
    ceny_om = 256 - tp.wavelength_to_pix(lmbda, tp2, 1)[1]
    # Transform the OM centroids onto the detector.
    cenx_d, ceny_d = ctd.rot_centroids(*rot_params, cenx_om, ceny_om,
                                       bound=False)

    # Create an interpolated 1D PSF at each required position.
    for i, vals in enumerate(zip(cenx_d, ceny_d, lmbda)):
        cenx, ceny, lbd = vals[0], vals[1], vals[2]
        # Evaluate the interpolation polynomials at the current wavelength.
        wb_i = np.polyval(coef_b, lbd)
        wr_i = np.polyval(coef_r, lbd)
        # Recenter the profile of both anchors on the correct Y-centroid.
        bax = np.arange(256+2*pad)-yd22+ceny
        Banch_i = np.interp(np.arange(256+2*pad), bax, Banch)
        rax = np.arange(256+2*pad)-yd25+ceny
        Ranch_i = np.interp(np.arange(256+2*pad), rax, Ranch)
        # Construct the interpolated profile.
        prof_int = (wb_i * Banch_i + wr_i * Ranch_i)
        # Re-add the lambda/D scaling.
        prof_int_cs = _chromescale(lbd, prof_int, ceny, invert=True)
        # Put the interpolated profile on the detector.
        map2D[:, int(round(cenx, 0))] = prof_int_cs

        # Note detector coordinates of the edges of the interpolated region.
        bend = int(round(cenx, 0))
        if i == 0:
            # 2.85µm (i=0) limit may be off the end of the detector.
            rend = np.max([int(round(cenx, 0)), 0])

    # Stitch together the interpolation and data.
    newmap = np.zeros((256+2*pad, 2048))
    # Insert interpolated data
    newmap[:, rend:bend] = map2D[:, rend:bend]
    # Bluer region is known from the CLEAR exposure.
    # Mask contamination from second and third orders and reconstruct wings.
    for col in range(bend, 2048):
        cens = [ycens['order 1'][1][col], ycens['order 2'][1][col],
                ycens['order 3'][1][col]]
        newmap[:, col] = reconstruct_wings(clear[:, col], ycens=cens, pad=pad)
    if F277 is not None:
        # Add on the F277W frame to the red of the model.
        # Reconstruct wing structure and pad.
        for col in range(rend):
            cens = [ycens['order 1'][1][col]]
            newmap[:, col] = reconstruct_wings(F277[:, col], ycens=cens,
                                               contamination=False, pad=pad)
    # Insert interpolated data to the red of the data.
    else:
        newmap[:, 0:rend] = map2D[:, 0:rend]

    # Column normalize.
    newmap /= np.nansum(newmap, axis=0)
    # Add noise floor to prevent arbitrarily low values in padded wings.
    floor = np.nanpercentile(newmap[pad:(-1-pad), :], 2)
    newmap += floor

    return newmap

def construct_order2(clear, ref_file_args, extract_params):
    '''Preforms an extraction and reconstructs the detector with only the
    first order trace profile. The second order profile is then isolated
    through the difference of the original and reconstructed detector.

    Parameters
    ----------
    clear : np.array of float (2D)
        CLEAR data frame.
    ref_file_args : list
        List of parameters of the reference files.
    extract_params : dict
        Dictionary of arguments required by the extraction algorithm.

    Returns
    -------
    residual : np.array of float (2D)
        Detector with only the order 2 trace profile.
    '''

    # Set up the extraction.
    extra = ExtractionEngine(*ref_file_args, **extract_params)
    # Preform the extraction with only the first order.
    f_k = extra.extract(data=clear)
    # Rebuild the detector.
    rebuilt = extra.rebuild(f_k)
    rebuilt[np.isnan(rebuilt)] = 0
    # Isolate the second order by subtracting the reconstructed first
    # order from the data
    residual = clear - rebuilt

    return residual

def get_extract_params():
    '''
    '''
    params = {}
    # Map of expected noise (sig)
    bkgd_noise = 20.
    # Oversampling
    params["n_os"] = 1
    # Threshold on the spatial profile
    params["thresh"] = 1e-6

    return params

def get_ref_file_args(o1frame):
    '''
    '''
    path = '/Users/michaelradica/Documents/GitHub/jwst-mtl/SOSS/'
    # List of orders to consider in the extraction
    order_list = [1]

    # Wavelength solution
    wave_maps = []
    wave_maps.append(fits.getdata(path+"extract/Ref_files/wavelengths_m1.fits"))

    # Spatial profiles
    spat_pros = []
    spat_pros.append(o1frame)

    # Convert data from fits files to float (fits precision is 1e-8)
    wave_maps = [wv.astype('float64') for wv in wave_maps]
    spat_pros = [p_ord.astype('float64') for p_ord in spat_pros]

    # Throughputs
    thrpt_list = [ThroughputSOSS(order) for order in order_list]

    # Convolution kernels
    ker_list = [WebbKer(wv_map) for wv_map in wave_maps]

    # Put all inputs from reference files in a list
    ref_file_args = [wave_maps, spat_pros, thrpt_list, ker_list]

    return ref_file_args

def loicpsf(wavelist=None, wfe_real=None, filepath='', save_to_disk=True,
            oversampling=10, pixel=128, verbose=True):
    '''Calls the WebbPSF package to create monochromatic PSFs for NIRISS
    SOSS observations and save them to disk.

    Parameters
    ----------
    wavelist : list
        List of wavelengths (in meters) for which to generate PSFs.
    wfe_real : int
        Index of wavefront realization to use for the PSF (if non-default
        WFE realization is desired).
    filepath : str
        Path to the directory to which the PSFs will be written.
        Defaults to the current directory.
    save_to_disk : bool
        Whether to save PSFs to disk, or return them from the function.
    oversampling : int
        Oversampling pixels scale for the PSF.
    pixel : int
        Width of the PSF in native pixels.
    verbose : bool
        Whether to print explanatory comments.

    Returns
    -------
    psf_list : list
        If save_to_disk is False, a list of the generated PSFs.
    '''

    psf_list = []

    if wavelist is None:
        # List of wavelengths to generate PSFs for
        wavelist = np.linspace(0.5, 5.2, 95) * 1e-6
    # Dimension of the PSF in native pixels
    pixel = pixel

    # Select the NIRISS instrument
    niriss = webbpsf.NIRISS()

    # Override the default minimum wavelength of 0.6 microns
    niriss.SHORT_WAVELENGTH_MIN = 0.5e-6
    # Set correct filter and pupil wheel components
    niriss.filter = 'CLEAR'
    niriss.pupil_mask = 'GR700XD'

    # Change the WFE realization if desired
    if wfe_real is not None:
        niriss.pupilopd = ('OPD_RevW_ote_for_NIRISS_predicted.fits.gz',
                           wfe_real)

    # Loop through all wavelengths to generate PSFs
    for wave in wavelist:
        if verbose is True:
            print('Calculating PSF at wavelength ',
                  round(wave/1e-6, 2), ' microns')
        psf = niriss.calc_psf(monochromatic=wave, fov_pixels=pixel,
                              oversample=oversampling, display=False)
        psf_list.append(psf)

        if save_to_disk is True:
            # Save psf realization to disk
            text = '{0:5f}'.format(wave*1e+6)
            psf.writeto(str(filepath)+'SOSS_os'+str(oversampling)+'_'+str(pixel)
                        + 'x'+str(pixel)+'_'+text+'_'+str(wfe_real)+'.fits',
                        overwrite=True)

        if save_to_disk is False:
            return psf_list

def mask_order(frame, xpix, ypix):
    '''
    Depreciated!
    Create a pixel mask to isolate only the detector pixels belonging to
     a specific diffraction order.

    Parameters
    ----------
    frame : np.array of float (2D)
        Science data frame.
    xpix : np.array
        Data x centroids for the desired order
    ypix : np.array
        Data y centroids for the desired order

    Returns
    -------
    O1frame : np.array of float (2D)
        The input data frame, with all pixels other than those within
        +/- 20 pixels of ypix masked.
    '''

    mask = np.zeros([256, 2048])
    xx = np.round(xpix, 0).astype(int)
    yy = np.round(ypix, 0).astype(int)
    xr = np.linspace(np.min(xx), np.max(xx), np.max(xx)+1).astype(int)

    # Set all pixels within the extent of the spatial profile to 1
    for xxx, yyy in zip(xr, yy):
        # Ensure that we stay on the detector
        p_max = np.min([yyy+20, 255])
        p_min = np.max([yyy-21, 0])
        mask[p_min:p_max, xxx] = 1

    # Mask pixels not in the order we want
    O1frame = (mask * frame)

    return O1frame

def pad_spectral_axis(frame, xcens, ycens, pad=0):
    '''Add padding to the spectral axis by interpolating the corresponding
    edge profile onto a set of extrapolated centroids.

    Parameters
    ----------
    frame : np.array (2D)
        Data frame.
    xcens : list
        X-coordinates of the trace centroids.
    ycens : list
        Y-coordinates of the trace centroids.
    pad : int
        Amount of padding to add along either end of the spectral axis (in
        pixels).

    Returns
    -------
    newframe : np.array (2D)
        Data frame with padding on the spectral axis.
    '''

    ylen = int(frame.shape[0])
    xlen = int(frame.shape[1])
    pp = np.polyfit(xcens, ycens, 5)
    xax_pad = np.arange(xlen+2*pad)-pad
    ycens_pad = np.polyval(pp, xax_pad)

    newframe = np.zeros((ylen, xlen+2*pad))
    newframe[:, (pad):(xlen+pad)] = frame

    for col in range(pad):
        yax = np.arange(ylen)
        newframe[:, col] = np.interp(yax+ycens[5]-ycens_pad[col], yax,
                                     frame[:, 5])

    for col in range(xlen+pad, xlen+2*pad):
        yax = np.arange(ylen)
        newframe[:, col] = np.interp(yax+ycens[-10]-ycens_pad[col], yax,
                                     frame[:, -10])

    return newframe

def reconstruct_wings(profile, ycens=None, contamination=True, pad=0,
                      doplot=False):
    '''Masks the second and third diffraction orders and reconstructs the
     underlying wing structure of the first order. Also adds padding in the
     spatial direction if required.

    Parameters
    ----------
    profile : np.array
        Spectral trace spatial profile.
    ycens : list
        Y-coordinates of the trace centroids. Must include all three
        diffraction orders if contamination is True, or only the first order if
        False.
    contamination : bool
        If True, profile has contamination from the second and third
        diffraction orders.
    pad : int
        Amount to pad each end of the spartial axis (in pixels).
    doplot : bool
        If True, does diagnostic plotting.

    Returns
    -------
    newprof : np.array
        Input spatial profile with reconstructed wings and padding.

    Raises
    ------
    ValueError
        If centroids are not provided for all three orders when contamination
        is set to True.
    '''

    # Convert Y-centroid positions to indices
    ycens = np.atleast_1d(ycens)
    ycens = np.round(ycens, 0).astype(int)
    if contamination is True and ycens.size != 3:
        raise ValueError('Centroids must be provided for first three orders if there is contamination.')

    # ====== Reconstruct left wing ======
    # Get the left wing of the trace profile in log space.
    prof_l = np.log10(profile[0:(ycens[0]-12)])
    # and corresponding axis.
    axis_l = np.arange(256)[0:(ycens[0]-12)]

    # === Outlier masking ===
    # Mask the cores of each order.
    for order, ycen in enumerate(ycens):
        if order == 0:
            start = ycen-25
            end = len(profile)
        elif order == 1:
            start = np.max([ycen-15, 0])
            end = np.max([ycen+15, 1])
        else:
            start = np.max([ycen-10, 0])
            end = np.max([ycen+10, 1])
        # Set core of each order to NaN.
        prof_l[start:end] = np.nan

    # Fit the unmasked part of the wing to determine the mean trend.
    inds = np.where(np.isfinite(prof_l))[0]
    pp = _robust_linefit(axis_l[inds], prof_l[inds], (0, 0))
    wing_mean = pp[0]+pp[1]*axis_l[inds]
    # Calculate the standard dev of unmasked points from the mean trend.
    stddev = np.sqrt(np.median((prof_l[inds] - wing_mean)**2))
    # Find all outliers that are >3-sigma deviant from the mean.
    inds2 = np.where(prof_l[inds] - wing_mean > 3*stddev)

    # === Wing fit ===
    # Get fresh left wing profile.
    prof_l2 = np.log10(profile[0:(ycens[0]-12)])
    # Mask second and third orders.
    if contamination is True:
        for order, ycen in enumerate(ycens):
            if order == 1:
                start = np.max([ycen-15, 0])
                end = np.max([ycen+15, 1])
            elif order == 2:
                start = np.max([ycen-10, 0])
                end = np.max([ycen+10, 1])
            # Set core of each order to NaN.
            prof_l2[start:end] = np.nan
    # Mask outliers
    prof_l2[inds[inds2]] = np.nan

    # Indices of all unmasked points in the left wing.
    inds3 = np.isfinite(prof_l2)
    # Fit with a 7th order polynomial.
    pp_l = np.polyfit(axis_l[inds3], prof_l2[inds3], 7)

    # ====== Reconstruct right wing ======
    # Get the profile for the right wing in log space.
    prof_r = np.log10(profile[(ycens[0]+12):])
    # and corresponding axis.
    axis_r = np.arange(256)[(ycens[0]+12):]
    # Fit with third order polynomial.
    pp_r = np.polyfit(axis_r, prof_r, 3)

    # ===== Stitching =====
    # Find pixel to stitch right wing fit.
    jjr = ycens[0]+15
    iir = np.where(axis_r == jjr)[0][0]
    # Pad the right axis.
    axis_r_pad = np.linspace(axis_r[0], axis_r[-1]+pad, len(axis_r)+pad)
    # Join right wing to old trace profile.
    newprof = np.concatenate([profile[:jjr], 10**np.polyval(pp_r, axis_r_pad)[iir:]])

    # Find pixel to stitch left wing fit.
    jjl = ycens[0]-15
    # Pad the left axis.
    axis_l_pad = np.linspace(axis_l[0]-pad, axis_l[-1], len(axis_l)+pad)
    iil = np.where(axis_l_pad == jjl)[0][0]
    # Join left wing to old trace profile.
    newprof = np.concatenate([10**np.polyval(pp_l, axis_l_pad)[:iil], newprof[jjl:]])

    if doplot is True:
        plotting._plot_wing_reconstruction(profile, ycens, axis_l, axis_l_pad,
                                           axis_r_pad, pp_l, pp_r, prof_l2,
                                           newprof)

    return newprof

def reconstruct_wings2(frame, ycen, pad_factor=1):
    '''
    Depreciated!
    Takes a reconstructed trace profile which has been truncated about the
    centroid and reconstructs the extended wing structure using an exponential
    profile. Also adds padding in the spatial direction by extending the
    exponetial fit.

    Parameters
    ----------
    frame : np.ndarray of float (2D)
        Empirical trace model.
    ycen : np.array of float
        Y-pixel centroid positions.
    pad_factor : int
        Multiplicative padding factor on the spatial axis. Defaults to 1 (no
        padding).

    Returns
    -------
    newframe : np.ndarray of float (2D)
        Trace model with reconstructed extended wing structure, and required
        padding.
    '''

    # Create new detector array and spatial axis taking into account padding.
    newframe = np.zeros(((pad_factor)*frame.shape[0], frame.shape[1]))
    fullax = np.arange(frame.shape[0])
    fullax_pad = np.arange((pad_factor)*frame.shape[0]) - (frame.shape[0]/2)*(pad_factor-1)

    # Loop over each column on the detector.
    for col in range(frame.shape[1]):

        # Temporary hack for NaN columns
        if np.any(np.isnan(frame[:, col])):
            continue

        # Get the centroid Y-position.
        cen = int(round(ycen[col], 0))

        # Extract the left wing
        start = np.max([0, cen-75])
        ax_l = np.arange(start, cen-9)
        lwing = np.log10(frame[start:(cen-9), col])
        # Find where the log is finite
        lwing_noi = lwing[np.isfinite(lwing)]
        ax_l_noi = ax_l[np.isfinite(lwing)]
        # Fit a first order polynomial to the finite value of the log wing.
        # Equivalent to fitting an exponential to the wing.
        pp_l = np.polyfit(ax_l_noi, lwing_noi, 1)

        # Locate pixels for stitching - where the left wing goes to zero.
        ii_l = np.where(np.isinf(lwing))[0][-1]
        # Location in full axis.
        jj_l = np.where(fullax == ax_l[ii_l])[0][0]
        # Location in padded axis.
        kk_l = np.where(fullax_pad == ax_l[ii_l])[0][0]

        # Extract the right wing
        end = np.min([cen+50, 255])
        ax_r = np.arange(cen+9, end)
        rwing = np.log10(frame[(cen+9):end, col])
        # Find where the log is finite
        rwing_noi = rwing[np.isfinite(rwing)]
        ax_r_noi = ax_r[np.isfinite(rwing)]
        # Fit a first order polynomial to the finite value of the log wing.
        # Equivalent to fitting an exponential to the wing.
        pp_r = np.polyfit(ax_r_noi, rwing_noi, 1)

        # Locate pixels for stitching - where the right wing goes to zero.
        ii_r = np.where(np.isinf(rwing))[0][0]
        jj_r = np.where(fullax == ax_r[ii_r])[0][0]
        kk_r = np.where(fullax_pad == ax_r[ii_r])[0][0]

        # Stitch the wings to the original profile, add padding if necessary.
        newcol = np.concatenate([10**(np.polyval(pp_l, fullax_pad[:kk_l])),
                                 frame[jj_l:jj_r, col],
                                 10**(np.polyval(pp_r, fullax_pad[kk_r:]))])

        try:
            # Find any remaining pixels where the profile is zero.
            inds = np.where(newcol == 0)[0][0]
            # If there are remainining zeros, replace with mean of neighbours.
            newcol[inds] = np.mean([newcol[inds-1], newcol[inds+1]])
        except IndexError:
            pass

        # Renormalize the column.
        newframe[:, col] = newcol / np.nansum(newcol)

    return newframe

def _robust_linefit(x, y, p0):
    '''Wrapper around scipy's least_squares line fitting routine implementing
    the Huber loss function - to be more resistant to outliers.

    Parameters
    ----------
    x : list
        Data describing dependant variable.
    y : list
        Data describing independant variable.
    p0 : tuple
        Initial guess straight line parameters.

    Returns
    -------
    res.x : list
        Best fitting parameters of a straight line.
    '''

    def line_res(p, x, y):
        '''Residuals from a straight line'''
        return p[0]+p[1]*x-y

    # Preform outlier resistant fitting.
    res = least_squares(line_res, p0, loss='huber', f_scale=0.1, args=(x, y))

    return res.x


# -*- coding:utf-8 -*-

""" Tensorflow Implementation of Variational Autoencoder for sentences.
	Paper: Generating Sentences from a Continuous Space
	(https://arxiv.org/abs/1511.06349)
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf

class VAE(object):
	def __init__(self, hparam, word_to_index, index_to_word):
		self.hparam = hparam
		self.word_to_index = word_to_index
		self.index_to_word = index_to_word

		self.seq_len = hparam.seq_len
		self.vocab_size = hparam.vocab_size
		self.embedding_dim = hparam.embedding_dim
		self.hidden_dim = hparam.hidden_dim
		self.latent_dim = hparam.latent_dim
		self.clip_norm = hparam.clip_norm
		self.num_layers = hparam.num_layers
		self.beam_width = hparam.beam_width
		self.rnn_type = hparam.rnn_type
		self.anneal_max = hparam.anneal_max
		self.anneal_bias = hparam.anneal_bias
		self.anneal_slope = hparam.anneal_slope

		self.SOS_ID = word_to_index["<sos>"]
		self.EOS_ID = word_to_index["<eos>"]
		self.UNK_ID = word_to_index["<unk>"]
		self.PAD_ID = word_to_index["<pad>"]

		self._build_global_helper()
		self._build_graph()
		self._build_summary()

		self.merged_summary = tf.summary.merge_all()
		self.init = tf.global_variables_initializer()

	def _build_global_helper(self):
		self.inputs = tf.placeholder(tf.int32, [None, self.seq_len])
		self.targets = tf.placeholder(tf.int32, [None, self.seq_len])

		self.enc_inputs = self.inputs
		self.dec_inputs = self.inputs
		self.dec_targets = self.targets
		self.enc_lengths = tf.count_nonzero(self.enc_inputs, 1, dtype=tf.int32)
		self.dec_lengths = self.enc_lengths
		self.batch_size = tf.shape(self.inputs)[0]

		self.lr = tf.Variable(0.0, trainable=False)
		self.new_lr = tf.placeholder(tf.float32, [])
		self.lr_update = tf.assign(self.lr, self.new_lr)

		self.global_step = tf.Variable(0, trainable=False)
		self.saver = tf.train.Saver()

	def _build_graph(self):
		self._build_forward_graph()
		self._build_backward_graph()

	def _build_forward_graph(self):
		with tf.variable_scope("encoder", reuse=None):
			embedding = tf.get_variable("embedding", [self.vocab_size, self.embedding_dim], dtype=tf.float32)
			self.embedded_enc_inputs = tf.nn.embedding_lookup(embedding, self.enc_inputs)
			self.embedded_dec_inputs = tf.nn.embedding_lookup(embedding, self.dec_inputs)

			enc_cell = self.create_rnn_cell(reuse=None)
			enc_init = enc_cell.zero_state(self.batch_size, tf.float32)
			enc_outputs, enc_states = tf.nn.dynamic_rnn(cell = enc_cell,
														initial_state = enc_init,
														inputs = self.embedded_enc_inputs,
														sequence_length = self.enc_lengths,
														dtype = tf.float32)
			if self.rnn_type == "gru":
				enc_vector = enc_states[-1]
			elif self.rnn_type == "block" or self.rnn_type == "basic":
				enc_vector = tf.concat([enc_states[-1].c, enc_states[-1].h], 1)

			self.enc_mean = tf.layers.dense(enc_vector, self.latent_dim, name="enc_mean")
			self.enc_logvar = tf.layers.dense(enc_vector, self.latent_dim, name="enc_logvar")
			epsilon = tf.truncated_normal(tf.shape(self.enc_mean))
			self.latent_sample = self.enc_mean + tf.exp(0.5 * self.enc_logvar) * epsilon

		self.training_rnn_outputs, self.training_logits = self.decoder_training(self.latent_sample)
		self.reconstructions = tf.argmax(self.training_logits, 2)
		self.predicted_ids = self.decoder_inference(self.latent_sample)

	def _build_backward_graph(self):
		self.nll_loss = self.nll_loss_fn()
		self.kl_w = self.kl_w_fn()
		self.kl_loss = self.kl_loss_fn()
		self.total_loss = self.nll_loss + self.kl_loss * self.kl_w

		params = tf.trainable_variables()
		gradients = tf.gradients(self.total_loss, params)
		clipped_gradients, _ = tf.clip_by_global_norm(gradients, self.clip_norm)
		self.train_op = tf.train.AdamOptimizer(self.lr).apply_gradients(
			zip(clipped_gradients, params), global_step=self.global_step)

	def create_rnn_cell(self, reuse):
		def create_single_cell(reuse):
			if self.rnn_type == "gru":
				cell = tf.contrib.rnn.GRUCell(self.hidden_dim,
										  	  kernel_initializer=tf.orthogonal_initializer(),
										  	  reuse=reuse)
			elif self.rnn_type == "block":
				cell = tf.contrib.rnn.LSTMBlockCell(self.hidden_dim, forget_bias=0.0, reuse=reuse)
			elif self.rnn_type == "basic":
				cell = tf.contrib.rnn.BasicLSTMCell(self.hidden_dim, forget_bias=0.0, state_is_tuple=True, reuse=reuse)
			else:
				ValueError("Unsupported rnn_type.")
			return cell
		cell = tf.contrib.rnn.MultiRNNCell([create_single_cell(reuse) for _ in range(self.num_layers)], state_is_tuple=True)
		return cell

	def _build_summary(self):
		tf.summary.scalar("total_loss", self.total_loss)
		tf.summary.scalar("nll_loss", self.nll_loss)
		tf.summary.scalar("kl_loss", self.kl_loss)
		tf.summary.scalar("kl_w", self.kl_w)
		# tf.summary.scalar("ratio", self.nll_loss / self.kl_loss)

	def decoder_training(self, latent_sample):
		with tf.variable_scope("encoder", reuse=True):
			embedding = tf.get_variable("embedding", [self.vocab_size, self.embedding_dim], dtype=tf.float32)
		with tf.variable_scope("decoder", reuse=None):
			if self.rnn_type == "gru":
				init_state = tuple([tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans0")] * self.num_layers)
			else:
				c = tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans1")
				h = tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans2")
				init_state = tuple([tf.contrib.rnn.LSTMStateTuple(c=c, h=h)] * self.num_layers)

			lin_proj = tf.layers.Dense(self.vocab_size, _scope='decoder/dense', _reuse=None)
			dec_cell = self.create_rnn_cell(reuse=None)
			helper = tf.contrib.seq2seq.TrainingHelper(inputs = self.embedded_dec_inputs,
													   sequence_length = self.dec_lengths)
			decoder = tf.contrib.seq2seq.BasicDecoder(cell = dec_cell,
													  helper = helper,
													  initial_state = init_state)
			decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(decoder = decoder)
		return decoder_output.rnn_output, lin_proj.apply(decoder_output.rnn_output)

	def decoder_inference(self, latent_sample):
		with tf.variable_scope("encoder", reuse=True):
			embedding = tf.get_variable("embedding", [self.vocab_size, self.embedding_dim], dtype=tf.float32)
		with tf.variable_scope("decoder", reuse=True):
			if self.rnn_type == "gru":
				init_state = tuple([tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans0")] * self.num_layers)
			else:
				c = tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans1")
				h = tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans2")
				init_state = tuple([tf.contrib.rnn.LSTMStateTuple(c=c, h=h)] * self.num_layers)

			lin_proj = tf.layers.Dense(self.vocab_size, _scope='decoder/dense', _reuse=True)
			dec_cell = self.create_rnn_cell(reuse=True)
			decoder = tf.contrib.seq2seq.BeamSearchDecoder(cell=dec_cell,
													   embedding=embedding,
													   start_tokens=tf.tile(tf.constant([self.SOS_ID], dtype=tf.int32), [self.batch_size]),
													   end_token = tf.constant(self.EOS_ID, dtype=tf.int32),
													   initial_state = tf.contrib.seq2seq.tile_batch(init_state, self.beam_width),
													   beam_width = self.beam_width,
													   output_layer = lin_proj)
			decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(decoder = decoder,
											maximum_iterations = 2 * tf.reduce_max(self.enc_lengths))
		return decoder_output.predicted_ids[:, :, 0]

	def nll_loss_fn(self):
		mask = tf.sequence_mask(self.dec_lengths, tf.reduce_max(self.dec_lengths), dtype=tf.float32)
		return tf.reduce_sum(tf.contrib.seq2seq.sequence_loss(
					logits = self.training_logits,
					targets = self.dec_targets,
					weights = mask,
					average_across_timesteps = False,
					average_across_batch = True))

	def kl_w_fn(self):
		return self.anneal_max * tf.sigmoid((self.anneal_slope) * (tf.cast(self.global_step, tf.float32) - tf.constant(self.anneal_bias / 2)))

	def kl_loss_fn(self):
		return 0.5 * tf.reduce_sum(tf.exp(self.enc_logvar) + tf.square(self.enc_mean) - 1 - self.enc_logvar) / tf.to_float(self.batch_size)

	def train_session(self, sess, inputs, targets):
		feed_dict = {self.inputs: inputs,
					 self.targets: targets}
		fetches = [self.train_op,
				   self.nll_loss,
				   self.kl_loss,
				   self.kl_w,
				   self.total_loss,
				   self.global_step,
				   self.merged_summary]
		_, nll_loss, kl_loss, kl_w, total_loss, step, summary = sess.run(fetches, feed_dict)
		return {"nll_loss": nll_loss,
				"kl_loss": kl_loss,
				"kl_w": kl_w,
				"total_loss": total_loss,
				"step": step,
				"summary": summary}

	def test_session(self, sess, inputs, targets):
		feed_dict = {self.inputs: inputs,
					 self.targets: targets}
		nll_loss, total_loss, kl_loss = sess.run([self.nll_loss, self.total_loss, self.kl_loss], feed_dict)
		return {"nll_loss": nll_loss,
				"total_loss": total_loss,
				"kl_loss": kl_loss}

	def random_generate_session(self, sess):
		feed_dict = {self.latent_sample: np.random.randn(1, self.latent_dim),
					 self.enc_lengths: [self.seq_len],
					 self.batch_size: 1}
		predicted_ids = sess.run(self.predicted_ids, feed_dict)[0]
		ret = 'Random Generation: %s' % ' '.join([self.index_to_word[index] for index in predicted_ids])
		return ret

	def reconstruct_session(self, sess, sentence):
		idx_list = [self.get_word_index(w) for w in sentence.lower().split()][:self.seq_len]
		idx_list = idx_list + [self.PAD_ID] * (self.seq_len - len(idx_list))
		ret1 = 'Original: %s' % ' '.join([self.index_to_word[index] for index in idx_list])

		predicted_ids = sess.run(self.reconstructions, {self.inputs: np.atleast_2d(idx_list)})[0]
		ret2 = 'Reconstr: %s' % ' '.join([self.index_to_word[index] for index in predicted_ids])
		return ret1, ret2

	def get_word_index(self, word):
		index = self.word_to_index[word]
		return index if index < self.vocab_size else self.UNK_ID

	def get_parameter_size(self):
		all_vars = tf.global_variables()
		total_count = 0
		for item in all_vars:
			if "Adam" in item.name:
				continue
			shape = item.get_shape().as_list()
			if len(shape) == 0:
				total_count += 1
			else:
				size =  1
				for val in shape:
					size *= val
				total_count += size
		return total_count


# Based on code from astropy
# Licensed under a 3-clause BSD style license

import warnings

import os
import ctypes
from functools import partial

import numpy as np
from numpy.ctypeslib import ndpointer, load_library

from .utils import KernelSizeError, has_even_axis, raise_even_kernel_exception

LIBRARY_PATH = os.path.dirname(__file__)

try:
    _convolve = load_library("_convolve", LIBRARY_PATH)
except Exception:
    raise ImportError("Convolution C extension is missing. Try re-building astropy.")

# The GIL is automatically released by default when calling functions imported
# from libraries loaded by ctypes.cdll.LoadLibrary(<path>)

# Declare prototypes
# Boundary None
_convolveNd_c = _convolve.convolveNd_c
_convolveNd_c.restype = None
_convolveNd_c.argtypes = [ndpointer(ctypes.c_double, flags={"C_CONTIGUOUS", "WRITEABLE"}),  # return array
                          ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"),  # input array
                          ctypes.c_uint,  # N dim
                          # size array for input and result unless
                          # embed_result_within_padded_region is False,
                          # in which case the result array is assumed to be
                          # input.shape - 2*(kernel.shape//2). Note: integer division.
                          ndpointer(ctypes.c_size_t, flags="C_CONTIGUOUS"),
                          ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"),  # kernel array
                          ndpointer(ctypes.c_size_t, flags="C_CONTIGUOUS"),  # size array for kernel
                          ctypes.c_bool,  # nan_interpolate
                          ctypes.c_bool,  # embed_result_within_padded_region
                          ctypes.c_uint]  # n_threads

# Disabling all doctests in this module until a better way of handling warnings
# in doctests can be determined
__doctest_skip__ = ['*']

BOUNDARY_OPTIONS = [None, 'fill', 'wrap', 'extend']

MAX_NORMALIZATION = 100

def _copy_input_if_needed(input, dtype=float, order='C', nan_treatment=None,
                          mask=None, fill_value=None):
    # strip quantity attributes
    if hasattr(input, 'unit'):
        input = input.value
    output = input
    # Copy input
    try:
        # Anything that's masked must be turned into NaNs for the interpolation.
        # This requires copying. A copy is also needed for nan_treatment == 'fill'
        # A copy prevents possible function side-effects of the input array.
        if nan_treatment == 'fill' or np.ma.is_masked(input) or mask is not None:
            if np.ma.is_masked(input):
                # ``np.ma.maskedarray.filled()`` returns a copy, however there
                # is no way to specify the return type or order etc. In addition
                # ``np.nan`` is a ``float`` and there is no conversion to an
                # ``int`` type. Therefore, a pre-fill copy is needed for non
                # ``float`` masked arrays. ``subok=True`` is needed to retain
                # ``np.ma.maskedarray.filled()``. ``copy=False`` allows the fill
                # to act as the copy if type and order are already correct.
                output = np.array(input, dtype=dtype, copy=False, order=order, subok=True)
                output = output.filled(fill_value)
            else:
                # Since we're making a copy, we might as well use `subok=False` to save,
                # what is probably, a negligible amount of memory.
                output = np.array(input, dtype=dtype, copy=True, order=order, subok=False)

            if mask is not None:
                # mask != 0 yields a bool mask for all ints/floats/bool
                output[mask != 0] = fill_value
        else:
            # The call below is synonymous with np.asanyarray(array, ftype=float, order='C')
            # The advantage of `subok=True` is that it won't copy when array is an ndarray subclass. If it
            # is and `subok=False` (default), then it will copy even if `copy=False`. This uses less memory
            # when ndarray subclasses are passed in.
            output = np.array(input, dtype=dtype, copy=False, order=order, subok=True)
    except (TypeError, ValueError) as e:
        raise TypeError('input should be a Numpy array or something '
                        'convertible into a float array', e)
    return output

def convolve(array, kernel, boundary='fill', fill_value=0.,
             nan_treatment='interpolate', normalize_kernel=True, mask=None,
             preserve_nan=False, normalization_zero_tol=1e-8):
    """
    Convolve an array with a kernel.

    This routine differs from `scipy.ndimage.convolve` because
    it includes a special treatment for ``NaN`` values. Rather than
    including ``NaN`` values in the array in the convolution calculation, which
    causes large ``NaN`` holes in the convolved array, ``NaN`` values are
    replaced with interpolated values using the kernel as an interpolation
    function.

    Parameters
    ----------
    array : `numpy.ndarray`
        The array to convolve. This should be a 1, 2, or 3-dimensional array
        or a list or a set of nested lists representing a 1, 2, or
        3-dimensional array.
    kernel : `numpy.ndarray`
        The convolution kernel. The number of dimensions should match those for
        the array, and the dimensions should be odd in all directions.  If a
        masked array, the masked values will be replaced by ``fill_value``.
    boundary : str, optional
        A flag indicating how to handle boundaries:
            * `None`
                Set the ``result`` values to zero where the kernel
                extends beyond the edge of the array.
            * 'fill'
                Set values outside the array boundary to ``fill_value`` (default).
            * 'wrap'
                Periodic boundary that wrap to the other side of ``array``.
            * 'extend'
                Set values outside the array to the nearest ``array``
                value.
    fill_value : float, optional
        The value to use outside the array when using ``boundary='fill'``
    normalize_kernel : bool, optional
        Whether to normalize the kernel to have a sum of one.
    nan_treatment : {'interpolate', 'fill'}
        interpolate will result in renormalization of the kernel at each
        position ignoring (pixels that are NaN in the image) in both the image
        and the kernel.
        'fill' will replace the NaN pixels with a fixed numerical value (default
        zero, see ``fill_value``) prior to convolution
        Note that if the kernel has a sum equal to zero, NaN interpolation
        is not possible and will raise an exception.
    preserve_nan : bool
        After performing convolution, should pixels that were originally NaN
        again become NaN?
    mask : `None` or `numpy.ndarray`
        A "mask" array.  Shape must match ``array``, and anything that is masked
        (i.e., not 0/`False`) will be set to NaN for the convolution.  If
        `None`, no masking will be performed unless ``array`` is a masked array.
        If ``mask`` is not `None` *and* ``array`` is a masked array, a pixel is
        masked of it is masked in either ``mask`` *or* ``array.mask``.
    normalization_zero_tol: float, optional
        The absolute tolerance on whether the kernel is different than zero.
        If the kernel sums to zero to within this precision, it cannot be
        normalized. Default is "1e-8".

    Returns
    -------
    result : `numpy.ndarray`
        An array with the same dimensions and as the input array,
        convolved with kernel.  The data type depends on the input
        array type.  If array is a floating point type, then the
        return array keeps the same data type, otherwise the type
        is ``numpy.float``.

    Notes
    -----
    For masked arrays, masked values are treated as NaNs.  The convolution
    is always done at ``numpy.float`` precision.
    """

    if boundary not in BOUNDARY_OPTIONS:
        raise ValueError("Invalid boundary option: must be one of {}"
                         .format(BOUNDARY_OPTIONS))

    if nan_treatment not in ('interpolate', 'fill'):
        raise ValueError("nan_treatment must be one of 'interpolate','fill'")

    # OpenMP support is disabled at the C src code level, changing this will have
    # no effect.
    n_threads = 1

    # Keep refs to originals
    passed_kernel = kernel
    passed_array = array

    # The C routines all need float type inputs (so, a particular
    # bit size, endianness, etc.).  So we have to convert, which also
    # has the effect of making copies so we don't modify the inputs.
    # After this, the variables we work with will be array_internal, and
    # kernel_internal.  However -- we do want to keep track of what type
    # the input array was so we can cast the result to that at the end
    # if it's a floating point type.  Don't bother with this for lists --
    # just always push those as float.
    # It is always necessary to make a copy of kernel (since it is modified),
    # but, if we just so happen to be lucky enough to have the input array
    # have exactly the desired type, we just alias to array_internal
    # Convert kernel to ndarray if not already

    # Copy or alias array to array_internal
    array_internal = _copy_input_if_needed(passed_array, dtype=float, order='C',
                                           nan_treatment=nan_treatment, mask=mask,
                                           fill_value=np.nan)
    array_dtype = getattr(passed_array, 'dtype', array_internal.dtype)
    # Copy or alias kernel to kernel_internal
    kernel_internal = _copy_input_if_needed(passed_kernel, dtype=float, order='C',
                                            nan_treatment=None, mask=None,
                                            fill_value=fill_value)

    # Make sure kernel has all odd axes
    if has_even_axis(kernel_internal):
        raise_even_kernel_exception()

    # -----------------------------------------------------------------------
    # From this point onwards refer only to ``array_internal`` and
    # ``kernel_internal``.
    # Assume both are base np.ndarrays and NOT subclasses e.g. NOT
    # ``Kernel`` nor ``np.ma.maskedarray`` classes.
    # -----------------------------------------------------------------------

    # Check dimensionality
    if array_internal.ndim == 0:
        raise Exception("cannot convolve 0-dimensional arrays")
    elif array_internal.ndim > 3:
        raise NotImplementedError('convolve only supports 1, 2, and 3-dimensional '
                                  'arrays at this time')
    elif array_internal.ndim != kernel_internal.ndim:
        raise Exception('array and kernel have differing number of '
                        'dimensions.')

    array_shape = np.array(array_internal.shape)
    kernel_shape = np.array(kernel_internal.shape)
    pad_width = kernel_shape//2

    # For boundary=None only the center space is convolved. All array indices within a
    # distance kernel.shape//2 from the edge are completely ignored (zeroed).
    # E.g. (1D list) only the indices len(kernel)//2 : len(array)-len(kernel)//2
    # are convolved. It is therefore not possible to use this method to convolve an
    # array by a kernel that is larger (see note below) than the array - as ALL pixels would be ignored
    # leaving an array of only zeros.
    # Note: For even kernels the correctness condition is array_shape > kernel_shape.
    # For odd kernels it is:
    # array_shape >= kernel_shape OR array_shape > kernel_shape-1 OR array_shape > 2*(kernel_shape//2).
    # Since the latter is equal to the former two for even lengths, the latter condition is complete.
    if boundary is None and not np.all(array_shape > 2*pad_width):
        raise KernelSizeError("for boundary=None all kernel axes must be smaller than array's - "
                              "use boundary in ['fill', 'extend', 'wrap'] instead.")

    # NaN interpolation significantly slows down the C convolution
    # computation. Since nan_treatment = 'interpolate', is the default
    # check whether it is even needed, if not, don't interpolate.
    # NB: np.isnan(array_internal.sum()) is faster than np.isnan(array_internal).any()
    nan_interpolate = (nan_treatment == 'interpolate') and np.isnan(array_internal.sum())

    # Check if kernel is normalizable
    if normalize_kernel or nan_interpolate:
        kernel_sum = kernel_internal.sum()
        kernel_sums_to_zero = np.isclose(kernel_sum, 0, atol=normalization_zero_tol)

        if kernel_sum < 1. / MAX_NORMALIZATION or kernel_sums_to_zero:
            raise ValueError("The kernel can't be normalized, because its sum is "
                             "close to zero. The sum of the given kernel is < {}"
                             .format(1. / MAX_NORMALIZATION))

    # Mark the NaN values so we can replace them later if interpolate_nan is
    # not set
    if preserve_nan or nan_treatment == 'fill':
        initially_nan = np.isnan(array_internal)
        if nan_treatment == 'fill':
            array_internal[initially_nan] = fill_value

    # Avoid any memory allocation within the C code. Allocate output array
    # here and pass through instead.
    result = np.zeros(array_internal.shape, dtype=float, order='C')

    embed_result_within_padded_region = True
    array_to_convolve = array_internal
    if boundary in ('fill', 'extend', 'wrap'):
        embed_result_within_padded_region = False
        if boundary == 'fill':
            # This method is faster than using numpy.pad(..., mode='constant')
            array_to_convolve = np.full(array_shape + 2*pad_width, fill_value=fill_value, dtype=float, order='C')
            # Use bounds [pad_width[0]:array_shape[0]+pad_width[0]] instead of [pad_width[0]:-pad_width[0]]
            # to account for when the kernel has size of 1 making pad_width = 0.
            if array_internal.ndim == 1:
                array_to_convolve[pad_width[0]:array_shape[0]+pad_width[0]] = array_internal
            elif array_internal.ndim == 2:
                array_to_convolve[pad_width[0]:array_shape[0]+pad_width[0],
                                  pad_width[1]:array_shape[1]+pad_width[1]] = array_internal
            else:
                array_to_convolve[pad_width[0]:array_shape[0]+pad_width[0],
                                  pad_width[1]:array_shape[1]+pad_width[1],
                                  pad_width[2]:array_shape[2]+pad_width[2]] = array_internal
        else:
            np_pad_mode_dict = {'fill': 'constant', 'extend': 'edge', 'wrap': 'wrap'}
            np_pad_mode = np_pad_mode_dict[boundary]
            pad_width = kernel_shape // 2

            if array_internal.ndim == 1:
                np_pad_width = (pad_width[0],)
            elif array_internal.ndim == 2:
                np_pad_width = ((pad_width[0],), (pad_width[1],))
            else:
                np_pad_width = ((pad_width[0],), (pad_width[1],), (pad_width[2],))

            array_to_convolve = np.pad(array_internal, pad_width=np_pad_width,
                                       mode=np_pad_mode)

    _convolveNd_c(result, array_to_convolve,
                  array_to_convolve.ndim,
                  np.array(array_to_convolve.shape, dtype=ctypes.c_size_t, order='C'),
                  kernel_internal,
                  np.array(kernel_shape, dtype=ctypes.c_size_t, order='C'),
                  nan_interpolate, embed_result_within_padded_region,
                  n_threads)

    # So far, normalization has only occured for nan_treatment == 'interpolate'
    # because this had to happen within the C extension so as to ignore
    # any NaNs
    if normalize_kernel:
        if not nan_interpolate:
            result /= kernel_sum
    elif nan_interpolate:
        result *= kernel_sum

    if nan_interpolate and not preserve_nan and np.isnan(result.sum()):
        warnings.warn("nan_treatment='interpolate', however, NaN values detected "
                      "post convolution. A contiguous region of NaN values, larger "
                      "than the kernel size, are present in the input array. "
                      "Increase the kernel size to avoid this.")

    if preserve_nan:
        result[initially_nan] = np.nan

    # Convert result to original data type
    if array_dtype.kind == 'f':
        # Try to preserve the input type if it's a floating point type
        # Avoid making another copy if possible
        try:
            return result.astype(array_dtype, copy=False)
        except TypeError:
            return result.astype(array_dtype)
    else:
        return result


#   Copyright 2019 PyFLOSIC developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
import time
import copy
import numpy as np
from ase import neighborlist as NL
from flosic_os import xyz_to_nuclei_fod,ase2pyscf
from ase import io, units, Atoms, Atom
from ase.utils import natural_cutoffs
from pyscf import gto, dft
from flosic_scf import FLOSIC
from ase.constraints import FixAtoms
import sys

#        >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
#        >>> grids = dft.gen_grid.Grids(mol)
#        >>> grids.level = 4
#        >>> grids.build()

# We want to have
# - a list of all neigbors of each atom
# - a list of FOD's that correspond to each atom
# - a list of all neigbors for each FOD'

# - different dft mesh's for each subgroup of a molecule

class ON(object):
    """Provides functionality to implement O(N) scaling FLOSIC"""
    def __init__(self, mol, fod, grid_level=9):
        super(ON, self).__init__()
        self.mol = mol
        self.nspin = len(fod)
        self._fod = fod   # list of positions, Angst
        self.grid_level = grid_level
        self.eps_cutoff = 1.0e-10
        self.is_init = True
        self.add_ghosts = False
        self.nshell = 2
        #print(self.grid_level)
        #sys.exit()

    def build(self):
        """build the required data structures (e.g. neigbor lists)"""
        # convert to Bohr as internal unit
        self.onatoms = list()
        for s in range(self.nspin):
            self.onatoms.append(dict())

        self.fod = list()
        self.nfod = [0,0]
        for s in range(self.nspin):
            self.fod.append( self._fod[s].copy())
            self.nfod[s] = self._fod[s].shape[0]
            #print np.array_str(self.fod[s],precision=4)

        # build atom object (make sure to enter ccors in Angst)
        acoord = self.mol.atom_coords()
        self.atoms = Atoms()
        for na in range(self.mol.natm):
            aa = Atom(symbol=self.mol.atom_symbol(na),
                position=acoord[na]*units.Bohr)
            self.atoms.extend(aa)

        cutoffs = natural_cutoffs(self.atoms)

        ##print cutoffs
        self.nl = NL.NeighborList(cutoffs,
            self_interaction=False, bothways=True)
        self.nl.update(self.atoms)

        ##print self.nl.get_neighbors(1)

        # cut out the H atoms from the distances
        #non_H_ids = list()
        #for atmid in range(self.mol.natm):
        #    if self.mol.atom_pure_symbol(atmid) == 'H': continue
        #    non_H_ids.append(atmid)
        #
        #non_H_pos = np.zeros((len(non_H_ids),3), dtype=np.float64)
        #for atmid,i in zip(non_H_ids,range(len(non_H_ids))):
        #    non_H_pos[i,:] = self.atoms.positions[atmid,:]

        #self.fod_atm = [[[-1,0.0,'C']*4],[[-2,0.0,'C']*4]]
        #print self.fod_atm[0]
        #self.fod[0].shape[0],self.fod[1].shape[0]
        self.fod_atm = list()
        # find out the nearest atom for each FOD and deciced
        # if its a core FOD or not
        for s in range(self.nspin):
            lfod_atm = list()
            for i in range(self.fod[s].shape[0]):
                fpos = self.fod[s][i]
                # distance to atoms
                dists = np.linalg.norm(fpos - self.atoms.positions,
                    axis=1)
                # do not use H-Atoms as nearest atoms
                ite = 0
                nearest_atom = np.argsort(dists)[ite]
                while self.mol.atom_pure_symbol(nearest_atom) == 'H':
                    ite += 1
                    nearest_atom = np.argsort(dists)[ite]

                #print self.fod_atm
                #print self.fod_atm[s]
                #self.fod_atm[s][i][0] = nearest_atom
                #self.fod_atm[s][i][1] = dists[nearest_atom]
                #print dists[nearest_atom]
                fodtype = 'V'
                if dists[nearest_atom] < 0.175:
                    fodtype = 'C'
                if dists[nearest_atom] < 1.0e-4:
                    fodtype = 'C1s'
                #self.fod_atm[s][i][2] = fodtype
                #print fodtype
                #sys.exit()
                lfod_atm.append(
                    (nearest_atom,dists[nearest_atom],fodtype))
            self.fod_atm.append(lfod_atm)
        #print len(self.fod_atm)

        # make a list of all fods that correspond to a specific atom
        self.atm_fod = list()
        for na in range(self.mol.natm):
            lfodlist = [[],[]]
            for s in range(self.nspin):
                for i in range(self.nfod[s]):
                    if self.fod_atm[s][i][0] == na:
                        #print "fod", i, "belongs to atom", na
                        lfodlist[s].append(i)
            self.atm_fod.append(lfodlist)
        #print self.atm_fod

        # prepare the ao slices
        self.fod_ao_slc = list()
        for s in range(self.nspin):
            lfod_slc = list()
            for i in range(self.nfod[s]):
                slc = self.get_nbas_slices(s,i)
                lfod_slc.append(slc)
            self.fod_ao_slc.append(lfod_slc)
        #print len(self.fod_ao_slc[0]), len(self.fod_ao_slc)

        # prepare the ongrids
        print("Generating O(N) meshes: nshell={}, grid level={}"\
            .format(self.nshell, self.grid_level))
        self.fod_onmsh = list()
        self.fod_onmol = list()
        for s in range(self.nspin):
            lfod_onmsh = list()
            lfod_onmol = list()
            for i in range(self.nfod[s]):
                omol, ogrid = self.get_grid(s,i,lfod_onmsh)
                lfod_onmsh.append(ogrid)
                lfod_onmol.append(omol)
            self.fod_onmsh.append(lfod_onmsh)
            self.fod_onmol.append(lfod_onmol)

        print("Generating O(N) FOD lists ...")
        #print self.fod_onmsh
        #print self.onatoms
        #sys.exit()

        # prepare a list that contains all fods that
        # correspond to a given fod
        self.fod_fod = [[],[]]
        for s in range(self.nspin):
            for i in range(self.nfod[s]):
                # get the atom numbers that corresponds to the fod
                #print self.fod_atm[s][i]
                #print self.atm_fod[s][i]
                atids = self.onatoms[s][i]
                lfod_fod = list()
                for atid in atids:
                    # now find out which of all the other fods
                    # correspond to these atoms
                    for j in range(self.nfod[s]):
                        if i == j: continue
                        atnr = self.fod_atm[s][j][0]
                        if atnr in atids:
                            lfod_fod.append(j)
                # remove duplicates
                lfod_fod = list(set(lfod_fod))
                #print(lfod_fod)
                #print i, len(lfod_fod)
                self.fod_fod[s].append(lfod_fod)
        #print self.fod_fod[0]
        #print self.fod_fod[1]
        #sys.exit()

        # finally, build a fod-grp list
        # this list orders all fod that have the same mesh
        self.fodlist2group()

    def update(self,s,fpos):
        '''update the class according to the given (new) fod positions [in Angst.]'''
        pass

    def fodlist2group(self):
        """docstring for fodlist2group"""
        print('generate fodgroup list')
        # pre-generate the data structure
        self.fodgrps = list()
        for s in range(self.nspin):
            self.fodgrps.append(list())

        #print(self.fodgrps)
        #print(len(self.fodgrps))

        #for j in range(self.nfod[0]):
        #    print(self.onatoms[0][j])
        #sys.exit()

        # self.onatoms hält zu jedem fod id die entsprechenden
        # atome, die das mesh definieren
        for s in range(self.nspin):
            ttable =  [False for i in range(self.nfod[s])]
            #print(ttable)
            grpcnt = 0
            for j in range(self.nfod[s]):
                if ttable[j]: continue
                self.fodgrps[s].append([j])
                jonatoms = self.onatoms[s][j]
                jonatoms.sort()
                for jj in range(j+1,self.nfod[s]):
                    if ttable[jj]: continue
                    cpmonatoms = self.onatoms[s][jj]
                    cpmonatoms.sort()
                    if jonatoms == cpmonatoms:
                        #print('fod {} and {} have same onatoms {}'.format(j, jj, jonatoms))
                        self.fodgrps[s][grpcnt].append(jj)
                        ttable[jj] = True
                grpcnt += 1
                #print(self.fodgrps[s])
                #sys.exit()
        return

    def get_grid(self,s,fodid,lmsh,level=None,verbose=False):
        """Generate a grid object that corresponds to the
            O(N) structure for the given FOD
        """
        print(' -> building Vxc-Grid for FOD {} ...'.format(fodid), flush=True)
        #mol.atom_pure_symbol
        if level == None:
            level = self.grid_level
        onatoms = self.onatoms[s][fodid]
        #print(onatoms)

        #print('onatoms:', onatoms)
        #sys.exit()
        fodtype = self.fod_atm[s][fodid][2]
        ongrid = None

        # check if we already have meshes available
        # with the same onatoms
        #self.mf.on.fod_onmsh[self.s][fgrp[0]]
        onatoms = list(sorted(onatoms))
        pmshid = -1
        for j in range(fodid):
            #qq = list(sorted(self.onatoms[s][j]))
            #print(j, onatoms, qq)
            #print(lmsh)
            #self.fod_onmsh
            #print(onatoms)
            #print(self.onatoms[s][j].sort())
            if onatoms == list(sorted(self.onatoms[s][j])):
                ## we already have a mesh generated
                ongrid = lmsh[j]
                pmshid = j
        #sys.exit()
        # build string to generate Mole object
        mstr = ''
        for na in onatoms:
            sym = self.mol.atom_pure_symbol(na)
            pos = self.atoms.positions[na]
            #print sym, pos
            mstr += "{0} {1:0.12f} {2:0.12f} {3:0.12f};".format(
                sym,pos[0],pos[1],pos[2])

        # add ghost atom for valence descriptors
        if (fodtype == 'V') and (self.add_ghosts):
            sym = 'ghost:H'
            pos = self.fod[s][fodid]
            #print sym, pos
            mstr += "{0} {1:0.12f} {2:0.12f} {3:0.12f};".format(
                sym,pos[0],pos[1],pos[2])

        print('     type {}, na {}:  atoms in msh {}'.format(fodtype, self.fod_atm[s][fodid][0], onatoms))

        # build a Mole object from subsystem
        #isinstance(o, str):
        #if type(self.mol.basis) == dict:
        #    b = self.mol.basis
        #    b['ghost'] = gto.basis.load('sto3g', 'H')
        #else:
        #    b = {'default':self.mol.basis, 'ghost': gto.basis.load('sto3g', 'H')}

        b = self.mol.basis

        #print(mstr)
        #print("{0} {1}".format(fodid,b))

        #sys.exit()

        try:
            onmol =  gto.M(atom=mstr,basis=b)
        except RuntimeError:
            onmol =  gto.M(atom=mstr, basis=b, spin=1)
        onmol.verbose=0
        onmol.max_memory=1000
        #print(onmol.atom)
        # build the meshes
        if ongrid is None:
            _mdft = dft.UKS(onmol)
            _mdft.max_cycle = 0
            _mdft.grids.level = level
            _mdft.kernel()
            ongrid = copy.copy(_mdft.grids)
        else:
            print('     (mesh was already generated for fod {})'.format(pmshid))
        #ongrid = dft.gen_grid.Grids(onmol)
        #ongrid.prune = dft.gen_grid.nwchem_prune
        #ongrid.level = level
        #ongrid.build()

        #print(level)
        #print(ongrid.coords.shape)
        #print onmol
        #sys.exit()

        return (onmol,ongrid)

    # now we want to find out the basis functions that are
    # really needed for each FO
    # - we check the nearest atom
    # - we find the neigbors to that atom
    # - we find the nbas indices for those atoms (start, stop)
    #   and store them in a list
    def get_nbas_slices(self,s,fodid,verbose=False):
        """docstring for get_nbas_slices"""
        slices = list()
        near_atm = self.fod_atm[s][fodid]
        onatoms = [near_atm[0]]
        if (near_atm[2] == 'C') or (near_atm[2] == 'C1s'):
            if self.nshell == 2:
                nei_atm = self.get_neighbors(near_atm[0], nshell=1)
                for na in nei_atm:
                    onatoms.append(na)

        if near_atm[2] == 'V':
            nei_atm = self.get_neighbors(near_atm[0], nshell=self.nshell)
            for na in nei_atm:
                onatoms.append(na)

        # remove duplicates
        onatoms = list(set(onatoms))

        if self.nshell == -1:
            onatoms = list(range(self.mol.natm))

        # update class onatoms dict
        self.onatoms[s][fodid] = onatoms

        all_slices=self.mol.aoslice_by_atom()
        #print all_slices
        onnbas = 0
        for na in onatoms:
            lslice=(all_slices[na,2],all_slices[na,3])
            slices.append(lslice)
            onnbas += lslice[1]-lslice[0]

        # print all_slices[-1][-1]
        pc = 100.0 - onnbas / (all_slices[-1,-1] / 100.0)

        if verbose:
            print("Sparsity: {0} / {1} ({2:0.2f} %)"
                .format(onnbas,all_slices[-1,-1], pc))

        return slices

        #sys.exit()

    def get_on_dm(self,s,fodid,dm,flo=None):
        """Mask out all basis functions that are not required
        return FLO's and dm were unused ebtries are set to zero

        if condense is set to True, return dm and FLO's with only
        the size needed for the O(N) method (e.g. reduced)

        """
        if dm.shape[-1] != self.mol.aoslice_by_atom()[-1,-1]:
            print("DM has wrong shape")
            print(self.mol.aoslice_by_atom()[-1,-1])
            sys.exit()

        # find out which indices in the density matrix
        # belong to a given fod
        slcs = self.fod_ao_slc[s][fodid]
        slc = list()
        for sl in slcs:
            slc += list(range(sl[0],sl[1]))

        dmout = dm.copy()
        #floout = flo.copy()

        # zero out the dm
        for i in range(dm.shape[-1]):
            if i in slc: continue
            for j in range(dm.shape[-1]):
                if j in slc: continue
                #print "set dm zero", i, j
                dmout[i,j] = 0.0

        nchk = np.sum(dmout)
        if np.isnan(nchk):
            print('Got it')
            sys.exit()

        return dmout

    def get_neighbors(self, atmid, nshell=1):
        """
        Return the indices of the neighbor atoms to a given atom id.

        Args:
            atmid : int
                id of the atom for which the neigbors shall be returned

        Kwargs:
            nshell : int
                number of shells around the given atom that should be
                considered as neigbors
                1 -> means include only 'nearest neigbors'
                2 -> include nearest and next-nearest neigbors
                3 -> (and larger not yet implemented!)

        Returns:
            List of integers which are the indices of the neigbor atoms
            to atmid.
        """
        #if (nshell < 1) or (nshell > 2):
        #    raise ValueError("nshell must be 1 or 2")
        if nshell == -1:
            ret = list(range(self.mol.natm))
            return ret

        if nshell == 1:
            ret = list(self.nl.get_neighbors(atmid)[0])
        else:
            base = list(self.nl.get_neighbors(atmid)[0])
            nn = []
            for aid in base:
                nn += list(self.nl.get_neighbors(aid)[0])
            base.extend(nn)

            # remove duplicates
            ret = list(dict.fromkeys(base))

        return ret

    def print_stats(self):
        '''Print out statistics'''
        print('    --- O(N) stats ---')
        print('{:>4} {:>4} : {:>7} {:>7} {:>12}'.format('spin','nfod', 'nmsh', 'nbas', 'onatoms'))
        _nmshs = 0
        _nbas = 0
        for s in range(self.nspin):
            for j in range(self.nfod[s]):
                onmsh = self.fod_onmsh[s][j]
                nonmsh = onmsh.coords.shape[0]
                onat = self.onatoms[s][j]
                nbas = self.fod_onmol[s][j].aoslice_by_atom()[-1,-1]
                _nbas += nbas
                _nmshs += nonmsh
                print('{:4d} {:4d} : {:7d} {:7d}    {}'.format(s,j,nonmsh,nbas, onat))
        print('    ----------------')
        _nbas = float(_nbas)/(np.sum(self.nfod))
        _tnbas = self.mol.aoslice_by_atom()[-1,-1]
        _nmshs = float(_nmshs)/(np.sum(self.nfod))
        #print(_tnbas)
        print('vsic parameters\n')
        print('nmsh: avg {:7d}'.format(int(_nmshs)))
        print('nbas: avg {:7d}'.format(int(_nbas)))
        print('spar: {:11.1f} %'.format(100.0 - _nbas/(_tnbas*0.01)))
        print('')
        print('Number of fod-groups with identical mesh: {}'.format(len(self.fodgrps[0])))
        print('    ----------------')

def C1s_FixPos(on, s):
    '''Returns an ASE constraint that fixes the 1s core descriptor positions

    Args:
        on : ON object
            the object that corresponds to the mol/fod structure

        s : spin index

    Returns:
        ase.constraint object if there are C1s to fix, None otherwise
    '''
    c1sidx = list()
    mol = on.mol
    for fodid in range(on.nfod[s]):
        _na = mol.atom_pure_symbol(on.fod_atm[s][fodid][0])
        _ft = on.fod_atm[s][fodid][2]
        if (_ft == 'C1s'):
            c1sidx.append(fodid)

    c1score = None
    if len(c1sidx) > 0:
        c1score = FixAtoms(indices=c1sidx)

    return c1score

if __name__ == '__main__':
    CH3SH = '''
    C -0.04795000 +1.14952000 +0.00000000
    S -0.04795000 -0.66486000 +0.00000000
    H +1.28308000 -0.82325000 +0.00000000
    H -1.09260000 +1.46143000 +0.00000000
    H +0.43225000 +1.55121000 +0.89226000
    H +0.43225000 +1.55121000 -0.89226000
    '''
    # this are the spin-up descriptors
    fod1 = Atoms([
    Atom('X', (-0.04795000, -0.66486000, +0.00000000)),
    Atom('X', (-0.04795000, +1.14952000, +0.00000000)),
    Atom('X', (-1.01954312, +1.27662578, +0.00065565)),
    Atom('X', (+1.01316012, -0.72796570, -0.00302478)),
    Atom('X', (+0.41874165, +1.34380502, +0.80870475)),
    Atom('X', (+0.42024357, +1.34411742, -0.81146545)),
    Atom('X', (-0.46764078, -0.98842277, -0.72314717)),
    Atom('X', (-0.46848962, -0.97040067, +0.72108036)),
    Atom('X', (+0.01320210, +0.30892333, +0.00444147)),
    Atom('X', (-0.28022018, -0.62888360, -0.03731204)),
    Atom('X', (+0.05389371, -0.57381853, +0.19630494)),
    Atom('X', (+0.09262866, -0.55485889, -0.15751914)),
    Atom('X', (-0.05807583, -0.90413106, -0.00104673))])

    # this are the spin-down descriptors
    fod2 = Atoms([
    Atom('He',( -0.04795000, -0.66486000, +0.0000000)),
    Atom('He',( -0.04795000, +1.14952000, +0.0000000)),
    Atom('He',( +1.12523084, -0.68699049, +0.0301970)),
    Atom('He',( +0.40996981, +1.33508869, +0.8089839)),
    Atom('He',( +0.40987059, +1.34148952, -0.8106910)),
    Atom('He',( -0.49563876, -0.99517303, +0.6829207)),
    Atom('He',( -0.49640020, -0.89986161, -0.6743094)),
    Atom('He',( +0.00073876, +0.28757089, -0.0298617)),
    Atom('He',( -1.03186573, +1.29783767, -0.0035536)),
    Atom('He',( +0.04235081, -0.54885843, +0.1924678)),
    Atom('He',( +0.07365725, -0.59150454, -0.1951675)),
    Atom('He',( -0.28422422, -0.61466396, -0.0087913)),
    Atom('He',( -0.02352948, -1.0425011 ,+0.01253239))])

    b = 'ccpvdz'
    spin = 0
    charge = 0

    mol = gto.M(atom=CH3SH,
                basis=b,
                spin=spin,
                charge=charge)

    grid_level  = 7
    mol.verbose = 4
    mol.max_memory = 2000
    mol.build()
    xc = 'LDA,PW'

    # quick dft calculation
    mdft = dft.UKS(mol)
    mdft.xc = xc
    mdft.kernel()

    # build O(N) stuff
    myon = ON(mol,[fod1.positions,fod2.positions], grid_level=grid_level)
    myon.nshell = 2
    myon.build()

    # enable ONMSH
    m = FLOSIC(mol,xc=xc,fod1=fod1,fod2=fod2,grid_level=grid_level, init_dm=mdft.make_rdm1())
    m.max_cycle = 40
    m.set_on(myon)
    m.conv_tol = 1e-5

    m.preopt = False
    m.preopt_start_cycle=0
    m.preopt_fix1s = True
    m.preopt_fmin = 0.005

    m.kernel()

    print(m.fod_gradients())
    print(m.get_fforces())


#! /usr/bin/env python

"""
Module with ADI algorithm (median psf subtraction).
Carlos A. Gomez / ULg
"""

from __future__ import division

__author__ = 'C. Gomez @ ULg'
__all__ = ['adi']

import numpy as np
from ..conf import time_ini, timing
from ..var import get_annulus, mask_circle
from ..preproc import cube_derotate, cube_collapse, check_PA_vector
from ..pca.pca_local import define_annuli

def adi(cube, angle_list, fwhm=4, radius_int=0, asize=2, delta_rot=1,
        mode='fullfr', nframes=4, collapse='median', full_output=False,
        verbose=True):
    """ Algorithm based on Marois et al. 2006 on Angular Differential Imaging.
    First the median frame is subtracted, then the median of the four closest
    frames taking into account the pa_threshold (field rotation).

    Parameters
    ----------
    cube : array_like, 3d
        Input cube.
    angle_list : array_like, 1d
        Corresponding parallactic angle for each frame.
    fwhm : float
        Known size of the FHWM in pixels to be used. Default is 4.
    radius_int : int, optional
        The radius of the innermost annulus. By default is 0, if >0 then the
        central circular area is discarded.
    asize : int, optional
        The size of the annuli, in FWHM. Default is 2.
    delta_rot : int, optional
        Factor for increasing the parallactic angle threshold, expressed in FWHM.
        Default is 1 (excludes 1 FHWM on each side of the considered frame).
    mode : {"fullfr","annular"}, str optional
        In "simple" mode only the median frame is subtracted, in "annular" mode
        also the 4 closest frames given a PA threshold (annulus-wise) are
        subtracted.
    nframes : even int optional
        Number of frames to be used for building the optimized reference PSF
        when working in annular mode.
    collapse : {'median', 'mean', 'sum', 'trimmean'}, str optional
        Sets the way of collapsing the frames for producing a final image.
    full_output: boolean, optional
        Whether to return the final median combined image only or with other
        intermediate arrays.
    verbose : {True, False}, bool optional
        If True prints to stdout intermediate info.

    Returns
    -------
    frame : array_like, 2d
        Median combination of the de-rotated cube.
    If full_output is True:
    cube_out : array_like, 3d
        The cube of residuals.
    cube_der : array_like, 3d
        The derotated cube of residuals.

    """
    def find_indices(angle_list, frame, thr, nframes):
        """ Returns the indices to be left in frames library for optimized ADI.
        To find a more pythonic way to do this!
        """
        n = angle_list.shape[0]
        index_prev = 0
        index_foll = frame
        for i in range(0, frame):
            if np.abs(angle_list[frame]-angle_list[i]) < thr:
                index_prev = i
                break
            else:
                index_prev += 1
        for k in range(frame, n):
            if np.abs(angle_list[k]-angle_list[frame]) > thr:
                index_foll = k
                break
            else:
                index_foll += 1

        window = int(nframes/2)
        ind1 = index_prev-window
        ind1 = max(ind1, 0)
        ind2 = index_prev
        ind3 = index_foll
        ind4 = index_foll+window
        ind4 = min(ind4, n)
        indices = np.array(range(ind1,ind2)+range(ind3,ind4))
        #print ind1, ind2, ind3, ind4, indices
        return indices

    #***************************************************************************
    array = cube

    if not array.ndim == 3:
        raise TypeError('Input array is not a cube or 3d array.')
    if not array.shape[0] == angle_list.shape[0]:
        raise TypeError('Input vector or parallactic angles has wrong length.')
    if not nframes%2==0:
        raise TypeError('nframes argument must be even value.')

    n, y, _ = array.shape

    if verbose:  start_time = time_ini()

    angle_list = check_PA_vector(angle_list)

    #***************************************************************************
    # The median frame (basic psf reference) is first subtracted from each frame.
    #***************************************************************************
    ref_psf = np.median(array, axis=0)
    array = array - ref_psf

    if mode=='fullfr':
        if radius_int>0:
            cube_out = mask_circle(array, radius_int)
        else:
            cube_out = array
        if verbose:  print 'Median psf reference subtracted'

    elif mode=='annular':
        annulus_width = int(asize * fwhm)                            # equal size for all annuli
        n_annuli = int(np.floor((y/2-radius_int)/annulus_width))
        if verbose:  print 'N annuli =', n_annuli, ', FWHM =', fwhm, '\n'
        #***********************************************************************
        # The annuli are built, and the corresponding PA thresholds for frame
        # rejection are calculated. The PA rejection is calculated at center of
        # the annulus.
        #***********************************************************************
        cube_out = np.zeros_like(array)
        for ann in range(n_annuli):
            pa_threshold,inner_radius,_= define_annuli(angle_list, ann, n_annuli,
                                                       fwhm, radius_int,
                                                       annulus_width, delta_rot,
                                                       verbose)

            indices = get_annulus(array[0], inner_radius, annulus_width,
                                  output_indices=True)
            yy = indices[0]
            xx = indices[1]

            matrix = array[:, yy, xx]    # shape [nframes x npx_annulus]

            #*******************************************************************
            # A second optimized psf reference is subtracted from each frame.
            # For each frame we find *nframes*, depending on the PA threshold,
            # to construct this optimized psf reference.
            #*******************************************************************
            for frame in range(n):
                if pa_threshold != 0:
                    indices_left = find_indices(angle_list, frame, pa_threshold,
                                                nframes)
                    matrix_disc = matrix[indices_left]
                else:
                    matrix_disc = matrix

                ref_psf_opt = np.median(matrix_disc, axis=0)
                curr_frame = matrix[frame]
                subtracted = curr_frame - ref_psf_opt
                cube_out[frame][yy, xx] = subtracted
        if verbose:  print 'Optimized median psf reference subtracted'

    else:
        raise RuntimeError('Mode not recognized')

    cube_der = cube_derotate(cube_out, angle_list)
    frame = cube_collapse(cube_der, mode=collapse)
    if verbose:
        print 'Done derotating and combining'
        timing(start_time)
    if full_output:
        return cube_out, cube_der, frame
    else:
        return frame