#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# 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.
#
# Authors: James D. McClain
#          Timothy Berkelbach <tim.berkelbach@gmail.com>
#

import time
import numpy
from functools import reduce

from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc import scf
from pyscf.cc import gccsd
from pyscf.pbc.mp.kmp2 import get_frozen_mask, get_nmo, get_nocc
from pyscf.pbc.cc import kintermediates as imdk
from pyscf.lib.parameters import LOOSE_ZERO_TOL, LARGE_DENOM
from pyscf.pbc.lib import kpts_helper

DEBUG = False

#
# FIXME: When linear dependence is found in KHF and handled by function
# pyscf.scf.addons.remove_linear_dep_, different k-point may have different
# number of orbitals.
#

#einsum = numpy.einsum
einsum = lib.einsum

def kernel(cc, eris, t1=None, t2=None, max_cycle=50, tol=1e-8, tolnormt=1e-6,
           max_memory=2000, verbose=logger.INFO):
    """Exactly the same as pyscf.cc.ccsd.kernel, which calls a
    *local* energy() function."""
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(cc.stdout, verbose)

    assert (isinstance(eris, gccsd._PhysicistsERIs))
    if t1 is None and t2 is None:
        t1, t2 = cc.init_amps(eris)[1:]
    elif t1 is None:
        nocc = cc.nocc
        nvir = cc.nmo - nocc
        nkpts = cc.nkpts
        t1 = numpy.zeros((nkpts, nocc, nvir), numpy.complex128)
    elif t2 is None:
        t2 = cc.init_amps(eris)[2]

    cput1 = cput0 = (time.clock(), time.time())
    nkpts, nocc, nvir = t1.shape
    eold = 0
    eccsd = 0

    if isinstance(cc.diis, lib.diis.DIIS):
        adiis = cc.diis
    elif cc.diis:
        adiis = lib.diis.DIIS(cc, cc.diis_file)
        adiis.space = cc.diis_space
    else:
        adiis = None

    conv = False
    for istep in range(max_cycle):
        t1new, t2new = cc.update_amps(t1, t2, eris, max_memory)
        normt = numpy.linalg.norm(t1new - t1) + numpy.linalg.norm(t2new - t2)
        if cc.iterative_damping < 1.0:
            alpha = cc.iterative_damping
            t1, t2 = (1-alpha)*t1 + alpha*t1new, (1-alpha)*t2 + alpha*t2new
        else:
            t1, t2 = t1new, t2new
        t1new = t2new = None

        t1, t2 = cc.run_diis(t1, t2, istep, normt, eccsd - eold, adiis)
        eold, eccsd = eccsd, energy(cc, t1, t2, eris)
        log.info('istep = %d  E(CCSD) = %.15g  dE = %.9g  norm(t1,t2) = %.6g', istep, eccsd, eccsd - eold, normt)
        cput1 = log.timer('CCSD iter', *cput1)
        if abs(eccsd - eold) < tol and normt < tolnormt:
            conv = True
            break
    log.timer('CCSD', *cput0)
    return conv, eccsd, t1, t2

def energy(cc, t1, t2, eris):
    nkpts, nocc, nvir = t1.shape
    fock = eris.fock
    eris_oovv = eris.oovv.copy()
    e = 0.0 + 0j
    for ki in range(nkpts):
        e += einsum('ia,ia', fock[ki, :nocc, nocc:], t1[ki, :, :])
    t1t1 = numpy.zeros(shape=t2.shape, dtype=t2.dtype)
    for ki in range(nkpts):
        ka = ki
        for kj in range(nkpts):
            #kb = kj
            t1t1[ki, kj, ka, :, :, :, :] = einsum('ia,jb->ijab', t1[ki, :, :], t1[kj, :, :])
    tau = t2 + 2 * t1t1
    e += 0.25 * numpy.dot(tau.flatten(), eris_oovv.flatten())
    e /= nkpts
    if abs(e.imag) > 1e-4:
        logger.warn(cc, 'Non-zero imaginary part found in KCCSD energy %s', e)
    return e.real

def update_amps(cc, t1, t2, eris, max_memory=2000):
    time0 = time.clock(), time.time()
    log = logger.Logger(cc.stdout, cc.verbose)
    nkpts, nocc, nvir = t1.shape
    fock = eris.fock

    fov = fock[:, :nocc, nocc:].copy()
    foo = fock[:, :nocc, :nocc].copy()
    fvv = fock[:, nocc:, nocc:].copy()

    tau = imdk.make_tau(cc, t2, t1, t1)

    Fvv = imdk.cc_Fvv(cc, t1, t2, eris)
    Foo = imdk.cc_Foo(cc, t1, t2, eris)
    Fov = imdk.cc_Fov(cc, t1, t2, eris)
    Woooo = imdk.cc_Woooo(cc, t1, t2, eris)
    Wvvvv = imdk.cc_Wvvvv(cc, t1, t2, eris)
    Wovvo = imdk.cc_Wovvo(cc, t1, t2, eris)

    # Move energy terms to the other side
    for k in range(nkpts):
        Fvv[k] -= numpy.diag(numpy.diag(fvv[k]))
        Foo[k] -= numpy.diag(numpy.diag(foo[k]))

    # Get the momentum conservation array
    # Note: chemist's notation for momentum conserving t2(ki,kj,ka,kb), even though
    # integrals are in physics notation
    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)

    eris_ovvo = numpy.zeros(shape=(nkpts, nkpts, nkpts, nocc, nvir, nvir, nocc), dtype=t2.dtype)
    eris_oovo = numpy.zeros(shape=(nkpts, nkpts, nkpts, nocc, nocc, nvir, nocc), dtype=t2.dtype)
    eris_vvvo = numpy.zeros(shape=(nkpts, nkpts, nkpts, nvir, nvir, nvir, nocc), dtype=t2.dtype)
    for km, kb, ke in kpts_helper.loop_kkk(nkpts):
        kj = kconserv[km, ke, kb]
        # <mb||je> -> -<mb||ej>
        eris_ovvo[km, kb, ke] = -eris.ovov[km, kb, kj].transpose(0, 1, 3, 2)
        # <mn||je> -> -<mn||ej>
        # let kb = kn as a dummy variable
        eris_oovo[km, kb, ke] = -eris.ooov[km, kb, kj].transpose(0, 1, 3, 2)
        # <ma||be> -> - <be||am>*
        # let kj = ka as a dummy variable
        kj = kconserv[km, ke, kb]
        eris_vvvo[ke, kj, kb] = -eris.ovvv[km, kb, ke].transpose(2, 3, 1, 0).conj()

    # T1 equation
    t1new = numpy.zeros(shape=t1.shape, dtype=t1.dtype)
    for ka in range(nkpts):
        ki = ka
        t1new[ka] += numpy.array(fov[ka, :, :]).conj()
        t1new[ka] += einsum('ie,ae->ia', t1[ka], Fvv[ka])
        t1new[ka] += -einsum('ma,mi->ia', t1[ka], Foo[ka])
        for km in range(nkpts):
            t1new[ka] += einsum('imae,me->ia', t2[ka, km, ka], Fov[km])
            t1new[ka] += -einsum('nf,naif->ia', t1[km], eris.ovov[km, ka, ki])
            for kn in range(nkpts):
                ke = kconserv[km, ki, kn]
                t1new[ka] += -0.5 * einsum('imef,maef->ia', t2[ki, km, ke], eris.ovvv[km, ka, ke])
                t1new[ka] += -0.5 * einsum('mnae,nmei->ia', t2[km, kn, ka], eris_oovo[kn, km, ke])

    # T2 equation
    t2new = numpy.array(eris.oovv).conj()
    for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
        # Chemist's notation for momentum conserving t2(ki,kj,ka,kb)
        kb = kconserv[ki, ka, kj]

        Ftmp = Fvv[kb] - 0.5 * einsum('mb,me->be', t1[kb], Fov[kb])
        tmp = einsum('ijae,be->ijab', t2[ki, kj, ka], Ftmp)
        t2new[ki, kj, ka] += tmp

        #t2new[ki,kj,kb] -= tmp.transpose(0,1,3,2)
        Ftmp = Fvv[ka] - 0.5 * einsum('ma,me->ae', t1[ka], Fov[ka])
        tmp = einsum('ijbe,ae->ijab', t2[ki, kj, kb], Ftmp)
        t2new[ki, kj, ka] -= tmp

        Ftmp = Foo[kj] + 0.5 * einsum('je,me->mj', t1[kj], Fov[kj])
        tmp = einsum('imab,mj->ijab', t2[ki, kj, ka], Ftmp)
        t2new[ki, kj, ka] -= tmp

        #t2new[kj,ki,ka] += tmp.transpose(1,0,2,3)
        Ftmp = Foo[ki] + 0.5 * einsum('ie,me->mi', t1[ki], Fov[ki])
        tmp = einsum('jmab,mi->ijab', t2[kj, ki, ka], Ftmp)
        t2new[ki, kj, ka] += tmp

        for km in range(nkpts):
            # Wminj
            #   - km - kn + ka + kb = 0
            # =>  kn = ka - km + kb
            kn = kconserv[ka, km, kb]
            t2new[ki, kj, ka] += 0.5 * einsum('mnab,mnij->ijab', tau[km, kn, ka], Woooo[km, kn, ki])
            ke = km
            t2new[ki, kj, ka] += 0.5 * einsum('ijef,abef->ijab', tau[ki, kj, ke], Wvvvv[ka, kb, ke])

            # Wmbej
            #     - km - kb + ke + kj = 0
            #  => ke = km - kj + kb
            ke = kconserv[km, kj, kb]
            tmp = einsum('imae,mbej->ijab', t2[ki, km, ka], Wovvo[km, kb, ke])
            #     - km - kb + ke + kj = 0
            # =>  ke = km - kj + kb
            #
            # t[i,e] => ki = ke
            # t[m,a] => km = ka
            if km == ka and ke == ki:
                tmp -= einsum('ie,ma,mbej->ijab', t1[ki], t1[km], eris_ovvo[km, kb, ke])
            t2new[ki, kj, ka] += tmp
            t2new[ki, kj, kb] -= tmp.transpose(0, 1, 3, 2)
            t2new[kj, ki, ka] -= tmp.transpose(1, 0, 2, 3)
            t2new[kj, ki, kb] += tmp.transpose(1, 0, 3, 2)

        ke = ki
        tmp = einsum('ie,abej->ijab', t1[ki], eris_vvvo[ka, kb, ke])
        t2new[ki, kj, ka] += tmp
        # P(ij) term
        ke = kj
        tmp = einsum('je,abei->ijab', t1[kj], eris_vvvo[ka, kb, ke])
        t2new[ki, kj, ka] -= tmp

        km = ka
        tmp = einsum('ma,mbij->ijab', t1[ka], eris.ovoo[km, kb, ki])
        t2new[ki, kj, ka] -= tmp
        # P(ab) term
        km = kb
        tmp = einsum('mb,maij->ijab', t1[kb], eris.ovoo[km, ka, ki])
        t2new[ki, kj, ka] += tmp

    eia = numpy.zeros(shape=(nocc, nvir), dtype=t1new.dtype)
    for ki in range(nkpts):
        eia = foo[ki].diagonal()[:, None] - fvv[ki].diagonal()[None, :]
        # When padding the occupied/virtual arrays, some fock elements will be zero
        idx = numpy.where(abs(eia) < LOOSE_ZERO_TOL)[0]
        eia[idx] = LARGE_DENOM

        t1new[ki] /= eia

    eijab = numpy.zeros(shape=(nocc, nocc, nvir, nvir), dtype=t2new.dtype)
    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
    for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
        kb = kconserv[ki, ka, kj]
        eijab = (foo[ki].diagonal()[:, None, None, None] + foo[kj].diagonal()[None, :, None, None] -
                 fvv[ka].diagonal()[None, None, :, None] - fvv[kb].diagonal()[None, None, None, :])
        # Due to padding; see above discussion concerning t1new in update_amps()
        idx = numpy.where(abs(eijab) < LOOSE_ZERO_TOL)[0]
        eijab[idx] = LARGE_DENOM

        t2new[ki, kj, ka] /= eijab

    time0 = log.timer_debug1('update t1 t2', *time0)

    return t1new, t2new

class GCCSD(gccsd.GCCSD):
    def __init__(self, mf, frozen=0, mo_coeff=None, mo_occ=None):
        assert (isinstance(mf, scf.khf.KSCF))
        if not isinstance(mf, scf.kghf.KGHF):
            mf = scf.addons.convert_to_ghf(mf)
        self.kpts = mf.kpts
        gccsd.GCCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)

    @property
    def nkpts(self):
        return len(self.kpts)

    get_nocc = get_nocc
    get_nmo = get_nmo
    get_frozen_mask = get_frozen_mask

    def dump_flags(self):
        logger.info(self, '\n')
        logger.info(self, '******** PBC CC flags ********')
        gccsd.GCCSD.dump_flags(self)
        return self

    def init_amps(self, eris):
        time0 = time.clock(), time.time()
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        t1 = numpy.zeros((nkpts, nocc, nvir), dtype=numpy.complex128)
        t2 = numpy.zeros((nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir), dtype=numpy.complex128)
        self.emp2 = 0
        foo = eris.fock[:, :nocc, :nocc].copy()
        fvv = eris.fock[:, nocc:, nocc:].copy()
        fov = eris.fock[:, :nocc, nocc:].copy()
        eris_oovv = eris.oovv.copy()
        eia = numpy.zeros((nocc, nvir))
        eijab = numpy.zeros((nocc, nocc, nvir, nvir))

        kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts)
        for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
            kb = kconserv[ki, ka, kj]
            eijab = (foo[ki].diagonal()[:, None, None, None] + foo[kj].diagonal()[None, :, None, None] -
                     fvv[ka].diagonal()[None, None, :, None] - fvv[kb].diagonal()[None, None, None, :])
            # Due to padding; see above discussion concerning t1new in update_amps()
            idx = numpy.where(abs(eijab) < LOOSE_ZERO_TOL)[0]
            eijab[idx] = LARGE_DENOM

            t2[ki, kj, ka] = eris_oovv[ki, kj, ka] / eijab

        t2 = numpy.conj(t2)
        self.emp2 = 0.25 * numpy.einsum('pqrijab,pqrijab', t2, eris_oovv).real
        self.emp2 /= nkpts

        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2.real)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, t1, t2

    def ccsd(self, t1=None, t2=None, eris=None, **kwargs):
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        self.eris = eris
        self.converged, self.e_corr, self.t1, self.t2 = \
                kernel(self, eris, t1, t2, max_cycle=self.max_cycle,
                       tol=self.conv_tol,
                       tolnormt=self.conv_tol_normt,
                       max_memory=self.max_memory, verbose=self.verbose)
        if self.converged:
            logger.info(self, 'CCSD converged')
        else:
            logger.info(self, 'CCSD not converge')
        if self._scf.e_tot == 0:
            logger.info(self, 'E_corr = %.16g', self.e_corr)
        else:
            logger.info(self, 'E(CCSD) = %.16g  E_corr = %.16g', self.e_corr + self._scf.e_tot, self.e_corr)
        return self.e_corr, self.t1, self.t2

    def ao2mo(self, mo_coeff=None):
        nkpts = self.nkpts
        nmo = self.nmo
        mem_incore = nkpts**3 * nmo**4 * 8 / 1e6
        mem_now = lib.current_memory()[0]

        if (mem_incore + mem_now < self.max_memory) or self.mol.incore_anyway:
            return _make_eris_incore(self, mo_coeff)
        else:
            raise NotImplementedError

    def update_amps(self, t1, t2, eris, max_memory=2000):
        return update_amps(self, t1, t2, eris, max_memory)

    def amplitudes_to_vector(self, t1, t2):
        return numpy.hstack((t1.ravel(), t2.ravel()))

    def vector_to_amplitudes(self, vec, nmo=None, nocc=None):
        if nocc is None: nocc = self.nocc
        if nmo is None: nmo = self.nmo
        nvir = nmo - nocc
        nkpts = self.nkpts
        nov = nkpts * nocc * nvir
        t1 = vec[:nov].reshape(nkpts, nocc, nvir)
        t2 = vec[nov:].reshape(nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir)
        return t1, t2

CCSD = GCCSD

def _make_eris_incore(cc, mo_coeff=None):
    log = logger.Logger(cc.stdout, cc.verbose)
    cput0 = (time.clock(), time.time())
    eris = gccsd._PhysicistsERIs()
    kpts = cc.kpts
    nkpts = cc.nkpts
    nocc = cc.nocc
    nmo = cc.nmo
    nvir = nmo - nocc
    eris.nocc = nocc

    #if any(nocc != numpy.count_nonzero(cc._scf.mo_occ[k] > 0) for k in range(nkpts)):
    #    raise NotImplementedError('Different occupancies found for different k-points')

    if mo_coeff is None:
        # If mo_coeff is not canonical orbital
        # TODO does this work for k-points? changed to conjugate.
        raise NotImplementedError
        mo_coeff = cc.mo_coeff
    nao = mo_coeff[0].shape[0]
    dtype = mo_coeff[0].dtype

    moidx = get_frozen_mask(cc)
    nocc_per_kpt = numpy.asarray(get_nocc(cc, per_kpoint=True))
    nmo_per_kpt  = numpy.asarray(get_nmo(cc, per_kpoint=True))

    padded_moidx = []
    for k in range(nkpts):
        kpt_nocc = nocc_per_kpt[k]
        kpt_nvir = nmo_per_kpt[k] - kpt_nocc
        kpt_padded_moidx = numpy.concatenate((numpy.ones(kpt_nocc, dtype=numpy.bool),
                                              numpy.zeros(nmo - kpt_nocc - kpt_nvir, dtype=numpy.bool),
                                              numpy.ones(kpt_nvir, dtype=numpy.bool)))
        padded_moidx.append(kpt_padded_moidx)

    eris.mo_coeff = []
    eris.orbspin = []
    # Generate the molecular orbital coefficients with the frozen orbitals masked.
    # Each MO is tagged with orbspin, a list of 0's and 1's that give the overall
    # spin of each MO.
    #
    # Here we will work with two index arrays; one is for our original (small) moidx
    # array while the next is for our new (large) padded array.
    for k in range(nkpts):
        kpt_moidx = moidx[k]
        kpt_padded_moidx = padded_moidx[k]

        mo = numpy.zeros((nao, nmo), dtype=dtype)
        mo[:, kpt_padded_moidx] = mo_coeff[k][:, kpt_moidx]
        if hasattr(mo_coeff[k], 'orbspin'):
            orbspin_dtype = mo_coeff[k].orbspin[kpt_moidx].dtype
            orbspin = numpy.zeros(nmo, dtype=orbspin_dtype)
            orbspin[kpt_padded_moidx] = mo_coeff[k].orbspin[kpt_moidx]
            mo = lib.tag_array(mo, orbspin=orbspin)
            eris.orbspin.append(orbspin)
        # FIXME: What if the user freezes all up spin orbitals in
        # an RHF calculation?  The number of electrons will still be
        # even.
        else:  # guess orbital spin - assumes an RHF calculation
            assert (numpy.count_nonzero(kpt_moidx) % 2 == 0)
            orbspin = numpy.zeros(mo.shape[1], dtype=int)
            orbspin[1::2] = 1
            mo = lib.tag_array(mo, orbspin=orbspin)
            eris.orbspin.append(orbspin)
        eris.mo_coeff.append(mo)

    # Re-make our fock MO matrix elements from density and fock AO
    dm = cc._scf.make_rdm1(cc.mo_coeff, cc.mo_occ)
    fockao = cc._scf.get_hcore() + cc._scf.get_veff(cc._scf.cell, dm)
    eris.fock = numpy.asarray([reduce(numpy.dot, (mo.T.conj(), fockao[k], mo)) for k, mo in enumerate(eris.mo_coeff)])

    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
    # The bottom nao//2 coefficients are down (up) spin while the top are up (down).
    # These are 'spin-less' quantities; spin-conservation will be added manually.
    so_coeff = [mo[:nao // 2] + mo[nao // 2:] for mo in eris.mo_coeff]

    eri = numpy.empty((nkpts, nkpts, nkpts, nmo, nmo, nmo, nmo), dtype=numpy.complex128)
    fao2mo = cc._scf.with_df.ao2mo
    for kp, kq, kr in kpts_helper.loop_kkk(nkpts):
        ks = kconserv[kp, kq, kr]
        eri_kpt = fao2mo(
            (so_coeff[kp], so_coeff[kq], so_coeff[kr], so_coeff[ks]), (kpts[kp], kpts[kq], kpts[kr], kpts[ks]),
            compact=False)
        eri_kpt[(eris.orbspin[kp][:, None] != eris.orbspin[kq]).ravel()] = 0
        eri_kpt[:, (eris.orbspin[kr][:, None] != eris.orbspin[ks]).ravel()] = 0
        eri_kpt = eri_kpt.reshape(nmo, nmo, nmo, nmo)
        eri[kp, kq, kr] = eri_kpt

    # Check some antisymmetrized properties of the integrals
    if DEBUG:
        check_antisymm_3412(cc, cc.kpts, eri)

    # Antisymmetrizing (pq|rs)-(ps|rq), where the latter integral is equal to
    # (rq|ps); done since we aren't tracking the kpoint of orbital 's'
    eri = eri - eri.transpose(2, 1, 0, 5, 4, 3, 6)
    # Chemist -> physics notation
    eri = eri.transpose(0, 2, 1, 3, 5, 4, 6)

    # Set the various integrals
    eris.dtype = eri.dtype
    eris.oooo = eri[:, :, :, :nocc, :nocc, :nocc, :nocc].copy() / nkpts
    eris.ooov = eri[:, :, :, :nocc, :nocc, :nocc, nocc:].copy() / nkpts
    eris.ovoo = eri[:, :, :, :nocc, nocc:, :nocc, :nocc].copy() / nkpts
    eris.oovv = eri[:, :, :, :nocc, :nocc, nocc:, nocc:].copy() / nkpts
    eris.ovov = eri[:, :, :, :nocc, nocc:, :nocc, nocc:].copy() / nkpts
    eris.ovvv = eri[:, :, :, :nocc, nocc:, nocc:, nocc:].copy() / nkpts
    eris.vvvv = eri[:, :, :, nocc:, nocc:, nocc:, nocc:].copy() / nkpts

    log.timer('CCSD integral transformation', *cput0)
    return eris

def check_antisymm_3412(cc, kpts, integrals):
    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
    nkpts = len(kpts)
    diff = 0.0
    for kp, kq, kr in kpts_helper.loop_kkk(nkpts):
        ks = kconserv[kp, kr, kq]
        for p in range(integrals.shape[3]):
            for q in range(integrals.shape[4]):
                for r in range(integrals.shape[5]):
                    for s in range(integrals.shape[6]):
                        pqrs = integrals[kp, kq, kr, p, q, r, s]
                        rspq = integrals[kq, kp, kr, q, p, r, s]
                        cdiff = numpy.linalg.norm(pqrs - rspq).real
                        if diff > 1e-5:
                            print("AS diff = %.15g" % cdiff, pqrs, rspq, kp, kq, kr, ks, p, q, r, s)
                        diff = max(diff, cdiff)
    print("antisymmetrization : max diff = %.15g" % diff)
    if diff > 1e-5:
        print("Energy cutoff (or cell.mesh) is not enough to converge AO integrals.")
    return diff

def check_antisymm_12(cc, kpts, integrals):
    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
    nkpts = len(kpts)
    diff = 0.0
    for kp, kq, kr in kpts_helper.loop_kkk(nkpts):
        ks = kconserv[kp, kr, kq]
        for p in range(integrals.shape[3]):
            for q in range(integrals.shape[4]):
                for r in range(integrals.shape[5]):
                    for s in range(integrals.shape[6]):
                        pqrs = integrals[kp, kq, kr, p, q, r, s]
                        qprs = integrals[kq, kp, kr, q, p, r, s]
                        cdiff = numpy.linalg.norm(pqrs + qprs).real
                        if diff > 1e-5:
                            print("AS diff = %.15g" % cdiff, pqrs, qprs, kp, kq, kr, ks, p, q, r, s)
                        diff = max(diff, cdiff)
    print("antisymmetrization : max diff = %.15g" % diff)
    if diff > 1e-5:
        print("Energy cutoff (or cell.mesh) is not enough to converge AO integrals.")

def check_antisymm_34(cc, kpts, integrals):
    kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts)
    nkpts = len(kpts)
    diff = 0.0
    for kp, kq, kr in kpts_helper.loop_kkk(nkpts):
        ks = kconserv[kp, kr, kq]
        for p in range(integrals.shape[3]):
            for q in range(integrals.shape[4]):
                for r in range(integrals.shape[5]):
                    for s in range(integrals.shape[6]):
                        pqrs = integrals[kp, kq, kr, p, q, r, s]
                        pqsr = integrals[kp, kq, ks, p, q, s, r]
                        cdiff = numpy.linalg.norm(pqrs + pqsr).real
                        if diff > 1e-5:
                            print("AS diff = %.15g" % cdiff, pqrs, pqsr, kp, kq, kr, ks, p, q, r, s)
                        diff = max(diff, cdiff)
    print("antisymmetrization : max diff = %.15g" % diff)
    if diff > 1e-5:
        print("Energy cutoff (or cell.mesh) is not enough to converge AO integrals.")


# -*- coding: utf-8 -*-
"""
CIE 2017 Colour Fidelity Index
==============================

Defines the *CIE 2017 Colour Fidelity Index* (CFI) computation objects:

- :class:`colour.quality.ColourRendering_Specification_CIE2017`
- :func:`colour.quality.colour_fidelity_index_CIE2017`

References
----------
-   :cite:`CIETC1-902017` : CIE TC 1-90. (2017). CIE 2017 colour fidelity index
    for accurate scientific use. CIE Central Bureau. ISBN:978-3-902842-61-9
"""

import numpy as np
import os
from collections import namedtuple

from colour.algebra import Extrapolator, euclidean_distance, linstep_function
from colour.appearance import XYZ_to_CIECAM02, VIEWING_CONDITIONS_CIECAM02
from colour.colorimetry import (
    MSDS_CMFS_STANDARD_OBSERVER, MultiSpectralDistributions, SpectralShape,
    SpectralDistribution, sd_to_XYZ, sd_blackbody, reshape_msds, sd_ones,
    sd_CIE_illuminant_D_series)
from colour.models import XYZ_to_UCS, UCS_to_uv, JMh_CIECAM02_to_CAM02UCS
from colour.temperature import uv_to_CCT_Ohno2013, CCT_to_xy_CIE_D
from colour.utilities import CACHE_REGISTRY, as_int, usage_warning

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013-2021 - Colour Developers'
__license__ = 'New BSD License - https://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-developers@colour-science.org'
__status__ = 'Production'

__all__ = [
    'SPECTRAL_SHAPE_CIE2017', 'RESOURCES_DIRECTORY_CIE2017',
    'TCS_ColorimetryData_CIE2017', 'ColourRendering_Specification_CIE2017',
    'colour_fidelity_index_CIE2017', 'load_TCS_CIE2017',
    'CCT_reference_illuminant', 'sd_reference_illuminant',
    'tcs_colorimetry_data', 'delta_E_to_R_f'
]

SPECTRAL_SHAPE_CIE2017 = SpectralShape(380, 780, 1)
"""
Spectral shape for *CIE 2017 Colour Fidelity Index* (CFI)
standard.

SPECTRAL_SHAPE_CIE2017 : SpectralShape
"""

RESOURCES_DIRECTORY_CIE2017 = os.path.join(
    os.path.dirname(__file__), 'datasets')
"""
*CIE 2017 Colour Fidelity Index* resources directory.

RESOURCES_DIRECTORY_CIE2017 : unicode
"""

_CACHE_TCS_CIE2017 = CACHE_REGISTRY.register_cache(
    '{0}._CACHE_TCS_CIE2017'.format(__name__))

class TCS_ColorimetryData_CIE2017(
        namedtuple('TCS_ColorimetryData_CIE2017',
                   ('name', 'XYZ', 'CAM', 'JMh', 'Jpapbp'))):
    """
    Defines the the class storing *test colour samples* colorimetry data.
    """

class ColourRendering_Specification_CIE2017(
        namedtuple('ColourRendering_Specification_CIE2017',
                   ('name', 'sd_reference', 'R_f', 'R_s', 'CCT', 'D_uv',
                    'colorimetry_data', 'delta_E_s'))):
    """
    Defines the *CIE 2017 Colour Fidelity Index* (CFI) colour quality
    specification.

    Parameters
    ----------
    name : unicode
        Name of the test spectral distribution.
    sd_reference : SpectralDistribution
        Spectral distribution of the reference illuminant.
    R_f : numeric
        *CIE 2017 Colour Fidelity Index* (CFI) :math:`R_f`.
    R_s : array_like
        Individual *colour fidelity indexes* data for each sample.
    CCT : numeric
        Correlated colour temperature :math:`T_{cp}`.
    D_uv : numeric
        Distance from the Planckian locus :math:`\\Delta_{uv}`.
    colorimetry_data : tuple
        Colorimetry data for the test and reference computations.
    delta_E_s : ndarray, (16,)
        Colour shifts of samples.
    """

def colour_fidelity_index_CIE2017(sd_test, additional_data=False):
    """
    Returns the *CIE 2017 Colour Fidelity Index* (CFI) :math:`R_f` of given
    spectral distribution.

    Parameters
    ----------
    sd_test : SpectralDistribution
        Test spectral distribution.
    additional_data : bool, optional
        Whether to output additional data.

    Returns
    -------
    numeric or ColourRendering_Specification_CIE2017
        *CIE 2017 Colour Fidelity Index* (CFI) :math:`R_f`.

    References
    ----------
    :cite:`CIETC1-902017`

    Examples
    --------
    >>> from colour.colorimetry import SDS_ILLUMINANTS
    >>> sd = SDS_ILLUMINANTS['FL2']
    >>> colour_fidelity_index_CIE2017(sd)  # doctest: +ELLIPSIS
    70.1208254...
    """

    if sd_test.shape.start > 380 or sd_test.shape.end < 780:
        usage_warning('Test spectral distribution shape does not span the'
                      'recommended 380-780nm range, missing values will be'
                      'filled with zeros!')

        # NOTE: "CIE 2017 Colour Fidelity Index" standard recommends filling
        # missing values with zeros.
        sd_test = sd_test.copy()
        sd_test.extrapolator = Extrapolator
        sd_test.extrapolator_kwargs = {
            'method': 'constant',
            'left': 0,
            'right': 0
        }

    if sd_test.shape.interval > 5:
        raise ValueError('Test spectral distribution interval is greater than'
                         '5nm which is the maximum recommended value '
                         'for computing the "CIE 2017 Colour Fidelity Index"!')

    shape = SpectralShape(SPECTRAL_SHAPE_CIE2017.start,
                          SPECTRAL_SHAPE_CIE2017.end, sd_test.shape.interval)

    CCT, D_uv = CCT_reference_illuminant(sd_test)
    sd_reference = sd_reference_illuminant(CCT, shape)

    # NOTE: All computations except CCT calculation use the
    # "CIE 1964 10 Degree Standard Observer".
    # pylint: disable=E1102
    cmfs_10 = reshape_msds(
        MSDS_CMFS_STANDARD_OBSERVER['CIE 1964 10 Degree Standard Observer'],
        shape)

    # pylint: disable=E1102
    sds_tcs = reshape_msds(load_TCS_CIE2017(shape), shape)

    test_tcs_colorimetry_data = tcs_colorimetry_data(sd_test, sds_tcs, cmfs_10)
    reference_tcs_colorimetry_data = tcs_colorimetry_data(
        sd_reference, sds_tcs, cmfs_10)

    delta_E_s = np.empty(len(sds_tcs.labels))
    for i, _delta_E in enumerate(delta_E_s):
        delta_E_s[i] = euclidean_distance(
            test_tcs_colorimetry_data[i].Jpapbp,
            reference_tcs_colorimetry_data[i].Jpapbp)

    R_s = delta_E_to_R_f(delta_E_s)
    R_f = delta_E_to_R_f(np.average(delta_E_s))

    if additional_data:
        return ColourRendering_Specification_CIE2017(
            sd_test.name, sd_reference, R_f, R_s, CCT, D_uv,
            (test_tcs_colorimetry_data, reference_tcs_colorimetry_data),
            delta_E_s)
    else:
        return R_f

def load_TCS_CIE2017(shape):
    """
    Loads the *CIE 2017 Test Colour Samples* dataset appropriate for the given
    spectral shape.

    The datasets are cached and won't be loaded again on subsequent calls to
    this definition.

    Parameters
    ----------
    shape : SpectralShape
        Spectral shape of the tested illuminant.

    Returns
    -------
    MultiSpectralDistributions
        *CIE 2017 Test Colour Samples* dataset.

    Examples
    --------
    >>> sds_tcs = load_TCS_CIE2017(SpectralShape(interval=5))
    >>> len(sds_tcs.labels)
    99
    """

    global _CACHE_TCS_CIE2017

    interval = shape.interval

    assert interval in (1, 5), (
        'Spectral shape interval must be either 1nm or 5nm!')

    filename = 'tcs_cfi2017_{0}_nm.csv.gz'.format(as_int(interval))

    if filename in _CACHE_TCS_CIE2017:
        return _CACHE_TCS_CIE2017[filename]

    data = np.genfromtxt(
        str(os.path.join(RESOURCES_DIRECTORY_CIE2017, filename)),
        delimiter=',')
    labels = ['TCS{0} (CIE 2017)'.format(i) for i in range(99)]

    tcs = MultiSpectralDistributions(data[:, 1:], data[:, 0], labels)

    _CACHE_TCS_CIE2017[filename] = tcs

    return tcs

def CCT_reference_illuminant(sd):
    """
    Computes the reference illuminant correlated colour temperature
    :math:`T_{cp}` and :math:`\\Delta_{uv}` for given test spectral
    distribution using *Ohno (2013)* method.

    Parameters
    ----------
    sd : SpectralDistribution
        Test spectral distribution.

    Returns
    -------
    ndarray
        Correlated colour temperature :math:`T_{cp}`, :math:`\\Delta_{uv}`.

    Examples
    --------
    >>> from colour import SDS_ILLUMINANTS
    >>> sd = SDS_ILLUMINANTS['FL2']
    >>> CCT_reference_illuminant(sd)  # doctest: +ELLIPSIS
    (4224.4697052..., 0.0017871...)
    """

    XYZ = sd_to_XYZ(sd)

    CCT, D_uv = uv_to_CCT_Ohno2013(UCS_to_uv(XYZ_to_UCS(XYZ)))

    return CCT, D_uv

def sd_reference_illuminant(CCT, shape):
    """
    Computes the reference illuminant for a given correlated colour temperature
    :math:`T_{cp}` for use in *CIE 2017 Colour Fidelity Index* (CFI)
    computation.

    Parameters
    ----------
    CCT : numeric
        Correlated colour temperature :math:`T_{cp}`.
    shape : SpectralShape
        Desired shape of the returned spectral distribution.

    Returns
    -------
    SpectralDistribution
        Reference illuminant for *CIE 2017 Colour Fidelity Index* (CFI)
        computation.

    Examples
    --------
    >>> from colour.utilities import numpy_print_options
    >>> with numpy_print_options(suppress=True):
    ...     sd_reference_illuminant(  # doctest: +ELLIPSIS
    ...         4224.469705295263300, SpectralShape(380, 780, 20))
    SpectralDistribution([[ 380.        ,    0.0034089...],
                          [ 400.        ,    0.0044208...],
                          [ 420.        ,    0.0053260...],
                          [ 440.        ,    0.0062857...],
                          [ 460.        ,    0.0072767...],
                          [ 480.        ,    0.0080207...],
                          [ 500.        ,    0.0086590...],
                          [ 520.        ,    0.0092242...],
                          [ 540.        ,    0.0097686...],
                          [ 560.        ,    0.0101444...],
                          [ 580.        ,    0.0104475...],
                          [ 600.        ,    0.0107642...],
                          [ 620.        ,    0.0110439...],
                          [ 640.        ,    0.0112535...],
                          [ 660.        ,    0.0113922...],
                          [ 680.        ,    0.0115185...],
                          [ 700.        ,    0.0113155...],
                          [ 720.        ,    0.0108192...],
                          [ 740.        ,    0.0111582...],
                          [ 760.        ,    0.0101299...],
                          [ 780.        ,    0.0105638...]],
                         interpolator=SpragueInterpolator,
                         interpolator_kwargs={},
                         extrapolator=Extrapolator,
                         extrapolator_kwargs={...})
    """

    if CCT <= 5000:
        sd_planckian = sd_blackbody(CCT, shape)

    if CCT >= 4000:
        xy = CCT_to_xy_CIE_D(CCT)
        sd_daylight = sd_CIE_illuminant_D_series(xy).align(shape)

    if CCT < 4000:
        sd_reference = sd_planckian
    elif 4000 <= CCT <= 5000:
        # Planckian and daylight illuminant must be normalised so that the
        # mixture isn't biased.
        sd_planckian /= sd_to_XYZ(sd_planckian)[1]
        sd_daylight /= sd_to_XYZ(sd_daylight)[1]

        # Mixture: 4200K should be 80% Planckian, 20% CIE Illuminant D Series.
        m = (CCT - 4000) / 1000
        values = linstep_function(m, sd_planckian.values, sd_daylight.values)
        name = ('{0}K Blackbody & CIE Illuminant D Series Mixture - {1:.1f}%'
                .format(as_int(CCT), 100 * m))
        sd_reference = SpectralDistribution(values, shape.range(), name=name)
    elif CCT > 5000:
        sd_reference = sd_daylight

    return sd_reference

def tcs_colorimetry_data(sd_irradiance, sds_tcs, cmfs):
    """
    Returns the *test colour samples* colorimetry data under given test light
    source or reference illuminant spectral distribution for the
    *CIE 2017 Colour Fidelity Index* (CFI) computations.

    Parameters
    ----------
    sd_irradiance : SpectralDistribution
        Test light source or reference illuminant spectral distribution, i.e.
        the irradiance emitter.
    sds_tcs : MultiSpectralDistributions
        *Test colour samples* spectral distributions.
    cmfs : XYZ_ColourMatchingFunctions
        Standard observer colour matching functions.

    Returns
    -------
    list
        *Test colour samples* colorimetry data under the given test light
        source or reference illuminant spectral distribution.

    Examples
    --------
    >>> delta_E_to_R_f(4.4410383190)  # doctest: +ELLIPSIS
    70.1208254...
    """

    XYZ_w = sd_to_XYZ(sd_ones(), cmfs, sd_irradiance)
    Y_b = 20
    L_A = 100
    surround = VIEWING_CONDITIONS_CIECAM02['Average']

    tcs_data = []
    for sd_tcs in sds_tcs.to_sds():
        XYZ = sd_to_XYZ(sd_tcs, cmfs, sd_irradiance)
        CAM = XYZ_to_CIECAM02(XYZ, XYZ_w, L_A, Y_b, surround, True)
        JMh = CAM.J, CAM.M, CAM.h
        Jpapbp = JMh_CIECAM02_to_CAM02UCS(JMh)

        tcs_data.append(
            TCS_ColorimetryData_CIE2017(sd_tcs.name, XYZ, CAM, JMh, Jpapbp))

    return tcs_data

def delta_E_to_R_f(delta_E):
    """
    Converts from colour-appearance difference to
    *CIE 2017 Colour Fidelity Index* (CFI) :math:`R_f` value.

    Parameters
    ----------
    delta_E : numeric
        Euclidean distance between two colours in *CAM02-UCS* colourspace.

    Returns
    -------
    float
        Corresponding *CIE 2017 Colour Fidelity Index* (CFI) :math:`R_f` value.
    """

    c_f = 6.73

    return 10 * np.log(np.exp((100 - c_f * delta_E) / 10) + 1)


from __future__ import division
import numpy as np
from scipy.linalg import solve
import sys
import time
from properties import *
from parse_file import *

class RealTime(object):
    """A RealTime object contains important parsed data from a Gaussian RealTime
    log file.

    Attributes:
        name:            A string containing primary filename
        logfile:         A string representing the Gaussian realtime log file
        electricDipole:  Object containing x, y, z electric dipole moments (au)
        magneticDipole:  Object containing x, y, z magnetic dipole moments (au)
        electricField:   Object containing x, y, z electric field strengths (au)
        magneticField:   Object containing x, y, z magnetic field strengths (au)
        iops:            Dict containing IOps for 512
        envelope:        Dict containing field parameters printed in logfile
        time:            Array containing time (au)
        energy:          Array containing energy (au)
        frequency:       Array containing frequencies from *time* (au)
        fourier:         Array containing fourier transformed signal (au)
        propertyarrays:  List containing names of properties stored as arrays.
        truncate:        Method to truncate propertyarrays to a given length
        mmut_restart:    Integer containing how often MMUT restarts
        au2fs:           Scalar constant to convert au to femtoseconds
    """

    def __init__(self, name, prog="GAUSSIAN"):
        """Return a RealTime object whose logfile is *logfile*."""
        # Initialize data
        self.name           = name
        self.prog           = prog

        if prog == "GAUSSIAN":
          self.logfile        = name+'.log'
          self.iops           = {'132':['0'],
                                 '134':['0'],
                                 '177':['0'],
                                 '136':['0'],
                                 '137':['0'],
                                 '138':['0'],
                                 '139':['0'],
                                 '140':['0'],
                                 '141':['0'],
                                 '142':['0'],
                                 '143':['0'],
                                 '144':['0']}
        elif prog == "CQ":
          self.fieldFile    = name + "_RealTime_AppliedField.csv"
          self.dipoleFile   = name + "_RealTime_Dipole.csv"
          self.mullikenFile = name + "_RealTime_Mulliken.csv"
          self.occAFile     = name + "_RealTime_OrbOcc_Alpha.csv"
          self.occBFile     = name + "_RealTime_OrbOcc_Beta.csv"

        else:
          pass # Should throw an error here

        self.envelope       = {}
        self.electricDipole = ElectricDipole()
        self.magneticDipole = MagneticDipole()
        self.electricField  = ElectricField()
        self.magneticField  = MagneticField()
        self.orthonorm      = None
        self.step_size      = None
        self.total_steps    = None
        self.time           = None
        self.energy         = None
        self.frequency      = None
        self.fourier        = None
        self.au2fs          = 0.0241888425
        # TODO May want to look at a better way of defining which attributes are
        # arrays instead of just hard-coding them in.
        self.propertyarrays = ['electricDipole',
                               'magneticDipole',
                               'electricField',
                               'magneticField',
                               'time',
                               #FIXME valid for H2+ Rabi ONLY
                               'HOMO',
                               'LUMO',
                               'energy']
        self.truncate       = truncate
        self.min_length     = None
        self.mmut_restart   = 10000000000 # e.g. never restart
        #FIXME: ONLY FOR H2+ RABI
        self.HOMO           = None
        self.LUMO           = None

        # Call parser
        parse_file(self)

        if prog == "GAUSSIAN":
          decode_iops(self)

          # Make all arrays consistent length
          clean_data(self)

    def pade_tx(self,dipole_direction='x',spectra='abs',damp_const=5500,
        num_pts=10000):
        # num_pts: number of points to sample for pade transformation

        if (spectra.lower() == 'abs') or (spectra.lower() == 'power'):
            if dipole_direction.lower() == 'x':
                dipole = self.electricDipole.x
#                kick_strength = self.electricField.x[0]
            elif dipole_direction.lower() == 'y':
                dipole = self.electricDipole.y
#                kick_strength = self.electricField.y[0]
            elif dipole_direction.lower() == 'z':
                dipole = self.electricDipole.z
#                kick_strength = self.electricField.z[0]
            else:
                print "Not a valid direction for the dipole! Try: x,y,z "
        elif spectra.lower() == 'ecd':
            if dipole_direction.lower() == 'x':
                dipole = self.magneticDipole.x
                kick_strength = self.electricField.x[0]
            elif dipole_direction.lower() == 'y':
                dipole = self.magneticDipole.y
                kick_strength = self.electricField.y[0]
            elif dipole_direction.lower() == 'z':
                dipole = self.magneticDipole.z
                kick_strength = self.electricField.z[0]
            else:
                print "Not a valid direction for the dipole! Try: x,y,z "
        else:
            print "Not a valid spectra choice"

#        if np.isclose(kick_strength,0.0):
#            if dipole_direction.lower() == 'x':
#                kick_strength = max(self.electricField.x)
#            elif dipole_direction.lower() == 'y':
#                kick_strength = max(self.electricField.y)
#            elif dipole_direction.lower() == 'z':
#                kick_strength = max(self.electricField.z)
#            if np.isclose(kick_strength,0.0):
#                print "Kick strength = 0. Make sure you FFT'd the correct direction"
#                sys.exit(0)
#            print "It looks like you are not perturbing the field at time = 0"
#            print "so we are taking the maximum of the electric field instead"
#            print "This may not be the functionality you want."

        # skip is integer to skip every n-th value
        # skip = 1 would not skip any values, but skip = 10 would only
        # consider every tenth value
        skip = 1
        dipole = dipole - dipole[0]
        dipole = dipole[::skip]
        damp = np.exp(-(self.time-self.time[0])/float(damp_const))
        damp = damp[::skip]
        dipole = dipole * damp

        timestep = skip*(self.time[2] - self.time[1])
        M = len(dipole)
        N = int(np.floor(M / 2))

        print "N = ", N
        if N > num_pts:
            N = num_pts
        print "Trimmed points to: ", N

        # G and d are (N-1) x (N-1)
        # d[k] = -dipole[N+k] for k in range(1,N)
        d = -dipole[N+1:2*N]

        # Old code, which works with regular Ax=b linear solver.
        # G[k,m] = dipole[N - m + k] for m,k in range(1,N)
        #G = dipole[N + np.arange(1,N)[:,None] - np.arange(1,N)]
        #b = solve(G,d,check_finite=False)

        # Toeplitz linear solver using Levinson recursion
        # Should be O(n^2), and seems to work well, but if you get strange
        # results you may want to switch to regular linear solver which is much
        # more stable.
        try:
            from scipy.linalg import toeplitz, solve_toeplitz
        except ImportError:
            print "You'll need SciPy version >= 0.17.0"

        # Instead, form G = (c,r) as toeplitz
        #c = dipole[N:2*N-1]
        #r = np.hstack((dipole[1],dipole[N-1:1:-1]))
        b = solve_toeplitz((dipole[N:2*N-1],\
            np.hstack((dipole[1],dipole[N-1:1:-1]))),d,check_finite=False)

        # Now make b Nx1 where b0 = 1
        b = np.hstack((1,b))

        # b[m]*dipole[k-m] for k in range(0,N), for m in range(k)
        a = np.dot(np.tril(toeplitz(dipole[0:N])),b)

        p = np.poly1d(a)
        q = np.poly1d(b)

        # If you want energies greater than 2*27.2114 eV, you'll need to change
        # the default frequency range to something greater.
        self.frequency = np.arange(0,2,0.000025)
        W = np.exp(-1j*self.frequency*timestep)

        fw_re = np.real(p(W)/q(W))
        fw_im = np.imag(p(W)/q(W))

        if np.any(np.isinf(self.frequency)) or np.any(np.isnan(self.frequency)):
            print "Check your dT: frequency contains NaNs and/or Infs!"
            sys.exit(0)

        if spectra.lower() == 'abs':
#            self.fourier = \
#                np.abs(self.frequency**2/(2.0*np.pi)*fw_im*fw_re)
            self.fourier = \
                np.abs(1.e0/(2.0*np.pi)*np.abs(p(W)/q(W))**2)
            np.savetxt('fftdata_cm-1_pade.txt', np.transpose([self.frequency/4.55633E-6,self.fourier]))
        elif spectra.lower() == 'ecd':
            self.fourier = \
                (17.32*fw_re)/(np.pi*kick_strength)
        elif spectra.lower() == 'power':
            self.fourier = \
                (self.frequency*(fw_re**2 + fw_im**2))/(np.pi*kick_strength)

    def fourier_tx(self,dipole_direction='x',spectra='abs',damp_const=150,
                    zero_pad=None,auto=False):
        """Return a set of frequencies and fourier transforms of a time
        dependent signal, e.g. return fourier transform of the x component of
        the time varying electric dipole"""
        from scipy.fftpack import fft, fftfreq
        # Choose which signal to FFT
        if spectra.lower() == 'abs':
            if dipole_direction.lower() == 'x':
                dipole = self.electricDipole.x
#                kick_strength = self.electricField.x[0]
            elif dipole_direction.lower() == 'y':
                dipole = self.electricDipole.y
#                kick_strength = self.electricField.y[0]
            elif dipole_direction.lower() == 'z':
                dipole = self.electricDipole.z
#                kick_strength = self.electricField.z[0]
            else:
                print "Not a valid direction for the dipole! Try: x,y,z "
        elif spectra.lower() == 'ecd':
            if dipole_direction.lower() == 'x':
                dipole = self.magneticDipole.x
                kick_strength = self.electricField.x[0]
            elif dipole_direction.lower() == 'y':
                dipole = self.magneticDipole.y
                kick_strength = self.electricField.y[0]
            elif dipole_direction.lower() == 'z':
                dipole = self.magneticDipole.z
                kick_strength = self.electricField.z[0]
            else:
                print "Not a valid direction for the dipole! Try: x,y,z "
        else:
            print "Not a valid spectra choice"

#        if np.isclose(kick_strength,0.0):
#            if dipole_direction.lower() == 'x':
#                kick_strength = max(self.electricField.x)
#            elif dipole_direction.lower() == 'y':
#                kick_strength = max(self.electricField.y)
#            elif dipole_direction.lower() == 'z':
#                kick_strength = max(self.electricField.z)
#            if np.isclose(kick_strength,0.0):
#                print "Kick strength = 0. Make sure you FFT'd the correct direction"
#                sys.exit(0)
#            print "It looks like you are not perturbing the field at time = 0"
#            print "so we are taking the maximum of the electric field instead"

        if auto:
            dt = self.time[2] - self.time[1]
            damp_const = self.time[-1]/10.0
            line_width = (2.0/damp_const)*27.2114
            #print "Damp const = ", damp_const
            if line_width > 2.0:
                print "Large line width: ", "{0:.3f}".format(line_width)," eV"
                print "Spectra not meaningful. Exiting..."
                sys.exit(0)
            else:
                print "Line width (eV) = ", "{0:.3f}".format(line_width)

            dipole = dipole - dipole[0]
            damp = np.exp(-(self.time-self.time[0])/float(damp_const))
            dipole = dipole * damp

            resolution = 0.025 #eV
            zero_pad   = int(np.floor((2.0*np.pi*27.2114)/(resolution*dt))\
                - len(self.time))
            if(zero_pad < 0.0):
                zero_pad = 0.0
            print "Number zeros = ", zero_pad

            zero = np.linspace(0,0,zero_pad)
            dipole = np.hstack((dipole,zero))

        else:
            dipole = dipole - dipole[0]
            damp = np.exp(-(self.time-self.time[0])/float(damp_const))
            dipole = dipole * damp

            if zero_pad:
                zero = np.linspace(0,0,zero_pad)
                dipole = np.hstack((dipole,zero))

        fw = fft(dipole)
        fw_re = np.real(fw)
        fw_im = np.imag(fw)

        n = len(fw_re)
        m = int(n / 2)
        timestep = self.time[2] - self.time[1]
        self.frequency = fftfreq(n,d=timestep)*2.0*np.pi
        if np.any(np.isinf(self.frequency)) or np.any(np.isnan(self.frequency)):
            print "Check your dT: frequency contains NaNs and/or Infs!"
            sys.exit(0)

        if spectra.lower() == 'abs':
            self.fourier = \
                np.abs(self.frequency**2/(2.0*np.pi)*np.abs(fw)**2)
            np.savetxt('fftdata_cm-1_fourier.txt', np.transpose([self.frequency/4.55633E-6,self.fourier]))

        elif spectra.lower() == 'ecd':
            self.fourier = \
                (17.32*fw_re)/(np.pi*kick_strength)

        # Grab positive values only
        self.frequency = self.frequency[1:m]
        self.fourier   = self.fourier[1:m]

    def test(self):
        self.check_energy()
        self.check_iops()
        pass

    def check_energy(self):
        dE = abs(max(self.energy) - min(self.energy))
        t_maxE = self.time[np.argmax(self.energy)]
        t_minE = self.time[np.argmin(self.energy)]
        print "Energy conserved to: ", "{0:.2e}".format(dE), " au"
        print "Max energy at time:  ", t_maxE, " au"
        print "Min energy at time:  ", t_minE, " au"

    def check_field(self,tol=1e-6):
        if self.envelope['Field']:
            print "External field:      ", self.envelope['Envelope']
            print "Ex field matches:    ", np.allclose(self.electricField.x,
                self.expected_field('Ex'),atol=tol)
            print "Ey field matches:    ", np.allclose(self.electricField.y,
                self.expected_field('Ey'),atol=tol)
            print "Ez field matches:    ", np.allclose(self.electricField.z,
                self.expected_field('Ez'),atol=tol)
           # print "Bx field matches: ", np.allclose(self.magneticField.x,
           #     self.expected_field('Bx'),atol=tol)
           # print "By field matches: ", np.allclose(self.magneticField.y,
           #     self.expected_field('By'),atol=tol)
           # print "Bz field matches: ", np.allclose(self.magneticField.z,
           #     self.expected_field('Bz'),atol=tol)
        else:
            print "No external field applied"

    def check_iops(self):
        """ Check internal consistency of some set iops and values printed out
        to the logfile, as well as some derived quantities"""
        # Check the step size
        if self.step_size == (self.time[2] - self.time[1]):
            if ((self.step_size == 0.05) \
                and (int(self.iops['134'][0]) == 0)) or\
               (self.step_size == float(self.iops['134'][0])*0.00001):
                print "Time step             [OK]: ", self.step_size, " au"
        else:
            print "Inconsistent time step: "
            print "  IOps:                  ", self.iops['134'][1]
            print "  logfile header showing ", self.step_size
            print "  logfile showing        ", self.time[2] - self.time[1]
        # Check the total propagation steps
        if ((self.total_steps == 15) \
           and (int(self.iops['132'][0]) == 0)) or\
            (self.total_steps == abs(int(self.iops['132'][0]))):
                print "Number MMUT steps     [OK]: ", self.total_steps, " steps"
        else:
            print "Inconsistent propagation time: "
            print "  IOps:                  ", self.iops['132'][1]
            print "  logfile header showing ", self.total_steps
        # Check if external field is indeed On or OFF
        if ((self.envelope['Field'] == False) and\
           (int(self.iops['138'][0]) == 0)):
            print "Field off:            [OK]"
        elif (self.envelope and int(self.iops['138'][0]) != 0):
            print "Field on:             [OK]"
            self.check_field()
        else:
            print "Inconsistency in field:"
            print "IOps:                     ", self.iops['138']

        # Check Orthonormalization
        if ((self.orthonorm == self.iops['136'][1])):
            print "Orthonormality        [OK]:", self.orthonorm
        else:
            print "Inconsistency in orthonormality"
            print "IOps:                      ", self.iops['136'][1]
            print "logfile showing:           ", self.iops['136'][1]

    def expected_field(self,component):
        Time  = self.time
        TOn   = self.envelope['TOn']
        TOff  = self.envelope['TOff']
        try:
            Omega = self.envelope['Frequency']
        except KeyError:
            Omega = 0.0
        try:
            Phase = self.envelope['Phase']
        except KeyError:
            Phase = 0.0
        OmegT = Omega*(Time - TOn) + Phase
        field = np.zeros_like(self.time)
        if self.envelope['Envelope'] == 'Constant':
            # Step function, depending on how TOn and TOff are defined
            idx = np.where((Time >= TOn) & (Time < TOff))
            # in GDV OmegT begins at TOn as well
            field[idx] = self.envelope[component]*np.cos(OmegT[idx])
        elif self.envelope['Envelope'] == 'Linear':
            TMax = (2.0*np.pi)/Omega
            # Linearly ramp off to zero
            idx = np.where((Time >= TOn) & (Time <= TOff) & \
                (Time > TOff-TMax))
            field[idx] = self.envelope[component]*\
                         ((TOff-Time[idx])/TMax)*np.cos(OmegT[idx])
            # Constant envelope
            idx = np.where((Time >= TOn) & (Time <= TOff) & \
                (Time > TOn+TMax) & (Time <= TOff-TMax))
            field[idx] = self.envelope[component]*np.cos(OmegT[idx])
            # Linearly ramp up to maximum in first cycle
            idx = np.where((Time >= TOn) & (Time <= TOff) & \
                (Time <= TOn+TMax))
            field[idx] = self.envelope[component]*\
                         ((Time[idx]-TOn)/TMax)*np.cos(OmegT[idx])
        elif self.envelope['Envelope'] == 'Gaussian':
            idx = np.where((Time >= TOn) & (Time < TOff))
            #FIXME: Sigma is hard-coded for testing...need to print it in the
            # output and then search for it during parsing.
            Sigma = 0.01
            TCntr = np.sqrt(np.log(1000.0))/Sigma
            field[idx] = self.envelope[component]*\
                         np.cos(OmegT[idx])*\
                         np.exp(-(Sigma*(Time[idx]-TCntr))**2)
        else:
            print "Not a valid field!"
            sys.exit(0)
        return field

if __name__ == '__main__':
    a = RealTime('test')
    import matplotlib.pyplot as plt
    plt.plot(a.time,a.electricDipole.z)
    plt.savefig('dipole.pdf')
    #plt.show()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Common Colour Models Utilities
==============================

Defines various colour models common utilities.

See Also
--------
`RGB Colourspaces IPython Notebook
<http://nbviewer.ipython.org/github/colour-science/colour-ipython/blob/master/notebooks/models/rgb.ipynb>`_  # noqa
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.models import RGB_COLOURSPACES, XYZ_to_RGB

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013 - 2014 - Colour Developers'
__license__ = 'New BSD License - http://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-science@googlegroups.com'
__status__ = 'Production'

__all__ = ['XYZ_to_sRGB']

def XYZ_to_sRGB(XYZ,
                illuminant=RGB_COLOURSPACES.get('sRGB').whitepoint,
                chromatic_adaptation_method='CAT02',
                transfer_function=True):
    """
    Converts from *CIE XYZ* colourspace to *sRGB* colourspace.

    Parameters
    ----------
    XYZ : array_like, (3,)
        *CIE XYZ* colourspace matrix.
    illuminant : array_like, optional
        Source illuminant chromaticity coordinates.
    chromatic_adaptation_method : unicode, optional
        ('XYZ Scaling', 'Bradford', 'Von Kries', 'Fairchild', 'CAT02')
        *Chromatic adaptation* method.
    transfer_function : bool, optional
        Apply *sRGB* *transfer function*.

    Returns
    -------
    ndarray, (3,)
        *sRGB* colour matrix.

    Notes
    -----
    -   Input *CIE XYZ* colourspace matrix is in domain [0, 1].

    Examples
    --------
    >>> XYZ = np.array([0.1180583421, 0.1034, 0.0515089229])
    >>> XYZ_to_sRGB(XYZ)  # doctest: +ELLIPSIS
    array([ 0.4822488...,  0.3165197...,  0.2207051...])
    """

    sRGB = RGB_COLOURSPACES.get('sRGB')
    return XYZ_to_RGB(XYZ,
                      illuminant,
                      sRGB.whitepoint,
                      sRGB.to_RGB,
                      chromatic_adaptation_method,
                      sRGB.transfer_function if transfer_function else None)


"""
Andrin Jenal, 2017
ETH Zurich
"""

import tensorflow as tf
import numpy as np

import nn_ops
from vgg import vgg

class VAE_DCGAN:

    def __init__(self, image_size, channels, z_size=256, learning_rate_enc=5e-4, learning_rate_dis=5e-4, learning_rate_dec=5e-4):
        # summaries
        self.merged_summary_op = None
        self.summary_writer = None

        self.image_size = image_size
        self.image_channels = channels
        self.z_size = z_size

        self.d_real = 0.0
        self.d_fake = 0.0

        self.learning_rate_enc = learning_rate_enc
        self.learning_rate_gen = learning_rate_dec
        self.learning_rate_dis = learning_rate_dis

        self.eps = 1e-8

        self.lr_discriminator = tf.placeholder(tf.float32, shape=[])
        self.lr_generator = tf.placeholder(tf.float32, shape=[])
        self.lr_encoder = tf.placeholder(tf.float32, shape=[])

        self.x = tf.placeholder(tf.float32, shape=(None, self.image_size, self.image_size, self.image_channels))
        self.batch_size = tf.shape(self.x)[0]

        self.z_p = tf.random_normal((self.batch_size, self.z_size), mean=0, stddev=1)
        self.eps = tf.random_normal((self.batch_size, self.z_size), mean=0, stddev=1)

        # preload vgg network for loss backpropagation
        self.vgg_weights, self.vgg_mean_pixel = vgg.load_net("datasets/imagenet-vgg-verydeep-19.mat")

        with tf.variable_scope("vae_dcgan_model"):
            tf.summary.histogram("x_values", self.x)

            with tf.variable_scope("encoder"):
                self.z_x_mean, self.z_x_log_sigma = self._encoder(self.x)

            with tf.variable_scope("generator"):
                self.z_x = tf.add(self.z_x_mean, tf.multiply(tf.exp(self.z_x_log_sigma), self.eps))
                tf.summary.histogram("z", self.z_x)
                tf.summary.histogram("z_mu", self.z_x_mean)
                tf.summary.histogram("z_sigma", tf.exp(self.z_x_log_sigma))

                self.x_tilde = self._generator(self.z_x)
                tf.summary.histogram("x_tilde_values", self.x_tilde)

            with tf.variable_scope("discriminator"):
                self.dis_x_tilde_p, self.l_x_tilde = self._discriminator(self.x_tilde)
                tf.summary.histogram("predicted_x_tilde_values", self.dis_x_tilde_p)

            with tf.variable_scope("generator", reuse=True):
                self.x_p = self._generator(self.z_p)
                tf.summary.histogram("x_p_values", self.x_tilde)

            with tf.variable_scope("discriminator", reuse=True):
                self.dis_x, self.l_x = self._discriminator(self.x)
                tf.summary.histogram("predicted_x_values", self.dis_x)

            with tf.variable_scope("discriminator", reuse=True):
                self.dis_x_p, _ = self._discriminator(self.x_p)
                tf.summary.histogram("predicted_x_p_values", self.dis_x_p)

            with tf.variable_scope("losses"):
                self.prior = self._kl_divergence()

                self.discriminator_loss = self._wasserstein_gradient_penalty_discriminator_loss()
                self.generator_loss = self._wasserstein_gradient_penalty_generator_loss()
                self.lth_layer_loss = self._lth_layer_loss()
                self.mse_loss = self._pixel_loss()
                #self.feature_loss = self._vgg_feature_loss()

                self.dissimilarity_loss = 0.5 * self.lth_layer_loss + 0.5 * self.mse_loss

                self.loss_encoder = self.prior + self.dissimilarity_loss
                self.loss_generator = self.dissimilarity_loss + self.generator_loss
                self.loss_discriminator = self.discriminator_loss

                train_variables = tf.trainable_variables()
                self.encoder_vars = [var for var in train_variables if "encoder" in var.name]
                self.discriminator_vars = [var for var in train_variables if "discriminator" in var.name]
                self.generator_vars = [var for var in train_variables if "generator" in var.name]

                with tf.name_scope("encoder_optimizer"):
                    self.e_optim = self._adam_optimizer(self.loss_encoder, self.encoder_vars, self.lr_encoder)

                with tf.name_scope("discriminator_optimizer"):
                    self.d_optim = self._rms_prop_optimizer(self.loss_discriminator, self.discriminator_vars, self.lr_discriminator)

                with tf.name_scope("generator_optimizer"):
                    self.g_optim = self._rms_prop_optimizer(self.loss_generator, self.generator_vars, self.lr_generator)

        # initialize saver
        self.saver = tf.train.Saver([v for v in tf.global_variables() if "vae_dcgan_model" in v.name])

        self._check_tensors()

    def _adam_optimizer(self, loss, loss_params, learning_rate, beta1=0.5):
        optimizer = tf.train.AdamOptimizer(learning_rate, beta1=beta1)
        grads = optimizer.compute_gradients(loss, var_list=loss_params)
        grads = nn_ops.clip_gradient_norms(grads, 50)
        train_optimizer = optimizer.apply_gradients(grads)
        grad_norms = self._l2_norms(grads)
        tf.summary.histogram("gradient_l2_norms", grad_norms)
        return train_optimizer

    def _rms_prop_optimizer(self, loss, loss_params, learning_rate):
        optimizer = tf.train.RMSPropOptimizer(learning_rate)
        grads = optimizer.compute_gradients(loss, var_list=loss_params)
        grads = nn_ops.clip_gradient_norms(grads, 50)
        train_optimizer = optimizer.apply_gradients(grads)
        grad_norms = self._l2_norms(grads)
        tf.summary.histogram("gradient_l2_norms", grad_norms)
        return train_optimizer

    def _l2_norms(self, gradients):
        return [tf.nn.l2_loss(g) for g, v in gradients if g is not None]

    def _encoder2(self, x):
        x = tf.reshape(x, [self.batch_size, self.image_size * self.image_size * self.image_channels])
        z_mean = nn_ops.linear_contrib(x, self.z_size, activation_fn=None, scope="fully_connected")
        z_log_sigma_sq = nn_ops.linear_contrib(x, self.z_size, activation_fn=None, scope="fully_connected")
        return z_mean, z_log_sigma_sq

    def _encoder(self, x):
        x = tf.reshape(x, [self.batch_size, self.image_size, self.image_size, self.image_channels])
        conv1 = nn_ops.conv2d_contrib(x, 64, kernel=5, stride=2, activation_fn=nn_ops.relu_batch_norm, scope="conv1")
        conv2 = nn_ops.conv2d_contrib(conv1, 128, kernel=3, stride=2, activation_fn=nn_ops.relu_batch_norm, scope="conv2")
        conv3 = nn_ops.conv2d_contrib(conv2, 256, kernel=3, stride=2, padding="VALID", activation_fn=nn_ops.relu_batch_norm, scope="conv3")
        flatten = nn_ops.flatten_contrib(conv3)
        z_mean = nn_ops.linear_contrib(flatten, self.z_size, activation_fn=nn_ops.relu_batch_norm, scope="fully_connected")
        z_log_sigma = nn_ops.linear_contrib(flatten, self.z_size, activation_fn=nn_ops.relu_batch_norm, scope="fully_connected")
        return z_mean, z_log_sigma

    def _discriminator(self, x):
        x = tf.reshape(x, [self.batch_size, self.image_size, self.image_size, self.image_channels])
        conv1 = nn_ops.conv2d_contrib(x, 64, kernel=5, stride=2, activation_fn=nn_ops.leaky_relu_batch_norm, scope="conv1")
        conv2 = nn_ops.conv2d_contrib(conv1, 128, kernel=3, stride=2, activation_fn=nn_ops.leaky_relu_batch_norm, scope="conv2")
        conv3 = nn_ops.conv2d_contrib(conv2, 256, kernel=3, stride=2, activation_fn=nn_ops.leaky_relu_batch_norm, scope="conv3")
        conv4 = nn_ops.conv2d_contrib(conv3, 256, kernel=3, stride=2, padding="VALID", activation_fn=nn_ops.leaky_relu_batch_norm, scope="conv4")
        conv4 = nn_ops.flatten_contrib(conv4)
        fc = nn_ops.linear_contrib(conv4, 512, activation_fn=nn_ops.leaky_relu_batch_norm, scope="fully_connected")
        predicted = nn_ops.linear_contrib(fc, 1, activation_fn=tf.nn.sigmoid, scope="prediction")
        net = {"conv1": conv1, "conv2": conv2, "conv3": conv3, "conv4": conv4, "fc": fc}
        return predicted, [net["conv1"], net["conv2"], net["fc"]]

    def _generator(self, z):
        fc = nn_ops.linear_contrib(z, 4 * 4 * 512, activation_fn=None)
        z = tf.reshape(fc, shape=(tf.shape(z)[0], 4, 4, 512))
        deconv1 = nn_ops.conv2d_transpose_contrib(z, 512, kernel=4, stride=2, activation_fn=nn_ops.relu_batch_norm, scope="upconv1")
        deconv2 = nn_ops.conv2d_transpose_contrib(deconv1, 256, kernel=5, stride=2, activation_fn=nn_ops.relu_batch_norm, scope="upconv2")
        deconv3 = nn_ops.conv2d_transpose_contrib(deconv2, 128, kernel=5, stride=2, activation_fn=nn_ops.relu_batch_norm, scope="upconv3")
        deconv4 = nn_ops.conv2d_transpose_contrib(deconv3, self.image_channels, kernel=5, stride=2, activation_fn=None, scope="upconv4")
        return tf.nn.tanh(deconv4)

    def _discriminator_loss(self, eps=1e-8):
        with tf.name_scope("logits_discriminator_loss"):
            dis_loss = tf.reduce_mean(-1.0 * tf.log(tf.clip_by_value(self.dis_x, eps, 1.0)) -
                                      tf.log(tf.clip_by_value(1.0 - self.dis_x_p, eps, 1.0)) -
                                      tf.log(tf.clip_by_value(1.0 - self.dis_x_tilde_p, eps, 1.0)))
            tf.summary.scalar("discriminator_loss_mean", dis_loss)
            return dis_loss

    def _generator_loss(self, eps=1e-8):
        with tf.name_scope("logits_generator_loss"):
            gen_loss = tf.reduce_mean(-1.0 * tf.log(tf.clip_by_value(self.dis_x_p, eps, 1.0)) -
                                      tf.log(tf.clip_by_value(self.dis_x_tilde_p, eps, 1.0)))
            tf.summary.scalar("generator_loss_mean", gen_loss)
            return gen_loss

    def _kl_divergence(self):
        with tf.name_scope("kl_divergence_loss"):
            KL = tf.reduce_sum((-self.z_x_log_sigma + 0.5 * (tf.exp(2.0 * self.z_x_log_sigma) + tf.square(self.z_x_mean)) - 0.5), axis=-1)
            KL_mean = tf.reduce_mean(KL)
            tf.summary.histogram("KL_divergence", KL)
            tf.summary.scalar("kl_divergence_mean", KL_mean)
            return KL_mean

    def _pixel_loss(self):
        with tf.name_scope("pixel_loss"):
            pixel_loss = tf.reduce_mean(tf.nn.l2_loss(self.x - self.x_tilde) / np.prod(self.x.get_shape().as_list()[1:]))
            pixel_loss_weighted = 1.0 * pixel_loss
            tf.summary.scalar("pixel_loss_mean", pixel_loss_weighted)
            return pixel_loss_weighted

    def _lth_layer_loss(self):
        with tf.name_scope("lth_layer_loss"):
            lth_layer_loss = 0
            for l1, l2 in zip(self.l_x, self.l_x_tilde):
                lth_layer_loss += tf.nn.l2_loss((l1 - l2)) / (self.image_size * self.image_size * self.image_channels)
            lth_layer_loss *= 1.0 / len(self.l_x)
            tf.summary.scalar("lth_layer_loss_mean", lth_layer_loss)
            return lth_layer_loss

    def _discriminator_binary_cross_entropy_loss(self):
        with tf.name_scope("cross_entropy_discriminator_loss"):
            d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.dis_x, labels=tf.ones_like(self.dis_x)))
            d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.dis_x_p, labels=tf.zeros_like(self.dis_x_p)))
            dis_loss = d_loss_real + d_loss_fake
            tf.summary.scalar("discriminator_loss_mean", dis_loss)
            return dis_loss

    def _generator_binary_cross_entropy_loss(self):
        with tf.name_scope("cross_entropy_generator_loss"):
            gen_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.dis_x_p, labels=tf.ones_like(self.dis_x_p)))
            tf.summary.scalar("generator_loss_mean", gen_loss)
            return gen_loss

    def _wasserstein_discriminator_loss(self):
        """ https://github.com/igul222/improved_wgan_training/blob/master/gan_mnist.py """
        with tf.name_scope("wasserstein_discriminator_loss"):
            dis_loss = -tf.reduce_mean(self.dis_x) + tf.reduce_mean(self.dis_x_p) + tf.reduce_mean(self.dis_x_tilde_p)
            tf.summary.scalar("discriminator_loss_mean", dis_loss)
            return dis_loss

    def _wasserstein_generator_loss(self):
        """ https://github.com/igul222/improved_wgan_training/blob/master/gan_mnist.py """
        with tf.name_scope("wasserstein_generator_loss"):
            gen_loss = -tf.reduce_mean(self.dis_x_p) - tf.reduce_mean(self.dis_x_tilde_p)
            tf.summary.scalar("generator_loss_mean", gen_loss)
            return gen_loss

    def _wasserstein_gradient_penalty_discriminator_loss(self):
        """ https://github.com/igul222/improved_wgan_training/blob/master/gan_mnist.py """
        with tf.name_scope("wasserstein_gradient_penalty_discriminator_loss"):
            dis_loss = -tf.reduce_mean(self.dis_x) + tf.reduce_mean(self.dis_x_p) + tf.reduce_mean(self.dis_x_tilde_p)
            x_p = tf.reshape(self.x_p, [self.batch_size, -1])
            x = tf.reshape(self.x, [self.batch_size, -1])
            differences = x_p - x
            interpolates = x + (tf.random_uniform([self.batch_size, 1], minval=0, maxval=1) * differences)
            dis_interpolates, _ = self._discriminator(interpolates)
            gradients = tf.gradients(dis_interpolates, [interpolates])[0]
            slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
            gradient_penalty = tf.reduce_mean((slopes - 1.) ** 2)
            dis_loss += 10 * gradient_penalty
            tf.summary.scalar("discriminator_loss_mean", dis_loss)
            return dis_loss

    def _wasserstein_gradient_penalty_generator_loss(self):
        """ https://github.com/igul222/improved_wgan_training/blob/master/gan_mnist.py """
        with tf.name_scope("wasserstein_gradient_penalty_generator_loss"):
            gen_loss = -tf.reduce_mean(self.dis_x_p) - tf.reduce_mean(self.dis_x_tilde_p)
            tf.summary.scalar("generator_loss_mean", gen_loss)
            return gen_loss

    def _vgg_feature_loss(self):
        with tf.name_scope("feature_loss_vgg"):
            feature_layers = ["relu3_3"]
            # [0.22591736, 0.77408264]
            # [0.09079630, 0.33333333, 0.57587037]
            # [0.04912966, 0.16162175, 0.33837825, 0.45087034]
            # [0.03139669, 0.09036695, 0.20000000, 0.30963305, 0.36860331]
            # [0.02222448, 0.05677295, 0.12367752, 0.20965582, 0.27656038, 0.31110886]
            # [0.01683352, 0.03891270, 0.08120526, 0.14285714, 0.20450902, 0.24680159, 0.26888077]
            # [0.01336953, 0.02844961, 0.05647934, 0.09969787, 0.15030213, 0.19352066, 0.22155039, 0.23663047]
            feature_weights = [1.0]
            feature_losses = []
            for ith, layer in enumerate(feature_layers):
                x_weights, _ = self._vgg_layer_weights(self.x, layer, self.batch_size, self.image_size, self.image_size, self.vgg_weights, self.vgg_mean_pixel)
                x_tilde_weights, x_tilde_size = self._vgg_layer_weights(self.x_tilde, layer, self.batch_size, self.image_size, self.image_size, self.vgg_weights, self.vgg_mean_pixel)
                feature_losses.append(feature_weights[ith] * (tf.nn.l2_loss(x_weights - x_tilde_weights) / x_tilde_size))
            feature_loss = tf.reduce_mean(tf.convert_to_tensor(feature_losses))
            feature_loss_weighted = feature_loss
            tf.summary.scalar("feature_loss_mean", feature_loss_weighted)
            return feature_loss_weighted

    def _vgg_layer_weights(self, input_images, layer_name, batch_size, image_height, image_width, vgg_weights, vgg_mean_pixel, pooling="avg"):
        if self.image_channels == 3:
            input_images = tf.reshape(input_images, shape=[batch_size, image_height, image_width, self.image_channels])
        else:
            input_images = tf.reshape(input_images, shape=[batch_size, image_height, image_width])
            input_images = tf.stack([input_images, input_images, input_images], axis=-1)
        input_images_mean = vgg.preprocess(input_images, vgg_mean_pixel)
        net_forward_images = vgg.net_preloaded(vgg_weights, input_images_mean, pooling)
        weights_size = np.prod(net_forward_images[layer_name].get_shape().as_list()[1:])
        return net_forward_images[layer_name], weights_size

    def _check_tensors(self):
        if tf.trainable_variables():
            for v in tf.trainable_variables():
                print("%s : %s" % (v.name, v.get_shape()))

    def _sigmoid(self, x, shift, mult):
        """
        Using this sigmoid to discourage one network overpowering the other
        """
        return 1 / (1 + np.exp(-(x + shift) * mult))

    def update_params(self, sess, input_tensor):
        e_current_lr = self.learning_rate_enc# * self._sigmoid(np.mean(self.d_real), -.5, 15)
        g_current_lr = self.learning_rate_gen# * self._sigmoid(np.mean(self.d_real), -.5, 15)
        d_current_lr = self.learning_rate_dis# * self._sigmoid(np.mean(self.d_fake), -.5, 15)

        _, _, _ = sess.run([self.d_optim, self.e_optim, self.g_optim], feed_dict={self.x: input_tensor,
                                                                                self.lr_encoder: e_current_lr,
                                                                                self.lr_generator: g_current_lr,
                                                                                self.lr_discriminator: d_current_lr})

        kl, d_loss, g_loss, lth_layer, d_real, d_fake = sess.run([self.prior, self.discriminator_loss, self.generator_loss, self.dissimilarity_loss, self.dis_x, self.dis_x_p], feed_dict={self.x: input_tensor})

        self.d_real = d_real
        self.d_fake = d_fake

        return kl, d_loss, g_loss, lth_layer, e_current_lr, g_current_lr, d_current_lr

    def generate_samples(self, sess, num_samples):
        z = np.random.normal(size=(num_samples, self.z_size))
        samples = sess.run(self.x_p, feed_dict={self.z_p: z})
        return np.array(samples)

    def initialize_summaries(self, sess, summary_directory):
        self.merged_summary_op = tf.summary.merge_all()
        self.summary_writer = tf.summary.FileWriter(summary_directory, sess.graph)

    def update_summaries(self, sess, x, epoch):
        if self.merged_summary_op is not None:
            summary = sess.run(self.merged_summary_op, feed_dict={self.x: x})
            self.summary_writer.add_summary(summary, global_step=epoch)
            print("updated summaries...")

    def restore_model(self, sess, checkpoint_file):
        self.saver.restore(sess, checkpoint_file)
        print("model restored from:", checkpoint_file)


import numpy as np
from scipy import optimize
from .phase_correlation import get_ambient_flow
from .objects import get_obj_extent

LARGE_NUM = 1000

def euclidean_dist(vec1, vec2):
    """ Computes euclidean distance. """
    vec1 = np.array(vec1)
    vec2 = np.array(vec2)
    dist = np.sqrt(sum((vec1-vec2)**2))
    return dist

def get_sizeChange(size1, size2):
    """ Returns change in size of an echo as the ratio of the larger size to
    the smaller, minus 1. """
    if (size1 < 5) and (size2 < 5):
        return 0
    elif size1 >= size2:
        return size1/size2 - 1
    else:
        return size2/size1 - 1

def find_objects(search_box, image2):
    """ Identifies objects found in the search region. """
    if not search_box['valid']:
        obj_found = np.array(-1)
    else:
        search_area = image2[search_box['x1']:search_box['x2'],
                             search_box['y1']:search_box['y2']]
        obj_found = np.unique(search_area)
    return obj_found

def shifts_disagree(shift1, shift2, record, thresh):
    """ Returns True if shift disparity greater than MAX_SHIFT_DISP
    parameter. """
    shift1 = shift1*record.grid_size[1:]
    shift2 = shift2*record.grid_size[1:]
    shift_disparity = euclidean_dist(shift1, shift2)
    return shift_disparity/record.interval.seconds > thresh

def clip_shift(shift, record, params):
    """ Clips shift according to MAX_FLOW_MAG paramter. """
    shift_meters = shift * record.grid_size[1:]
    shift_mag = np.linalg.norm(shift_meters)
    velocity = shift_mag / record.interval.seconds
    if velocity > params['MAX_FLOW_MAG'] and shift_mag != 0:
        unit = shift_meters / shift_mag
        clipped = unit * params['MAX_FLOW_MAG'] * record.interval.seconds
        clipped_pix = clipped/record.grid_size[1:]
        return clipped_pix
    else:
        return shift

def correct_shift(
        local_shift, current_objects, obj_id1, global_shift, record,
        params):
    """ Takes in flow vector based on local phase correlation (see
    get_std_flow) and compares it to the last headings of the object and
    the global_shift vector for that timestep. Corrects accordingly.
    Note: At the time of this function call, current_objects has not yet been
    updated for the current frame and frame_new, so the id2s in current_objects
    correspond to the objects in the current frame. """
    global_shift = clip_shift(global_shift, record, params)

    # Note last_heads is defined using object centers! These jump around a lot
    # when tracking large objects and should therefore probably not be used
    # when tracking MCS systems!

    if current_objects is None:
        last_heads = None
    else:
        obj_index = current_objects['id2'] == obj_id1
        last_heads = current_objects['last_heads'][obj_index].flatten()
        last_heads = np.round(last_heads * record.interval_ratio, 2)
        if len(last_heads) == 0:
            last_heads = None

    if last_heads is None:
        if shifts_disagree(
                local_shift, global_shift, record, params['MAX_SHIFT_DISP']):
            case = 0
            corrected_shift = global_shift
        else:
            case = 1
            corrected_shift = (local_shift + global_shift)/2

    elif shifts_disagree(
            local_shift, last_heads, record, params['MAX_SHIFT_DISP']):
        if shifts_disagree(
                local_shift, global_shift, record, params['MAX_SHIFT_DISP']):
            case = 2
            corrected_shift = last_heads
        else:
            case = 3
            corrected_shift = local_shift

    else:
        case = 4
        # corrected_shift = (local_shift + last_heads)/2
        if shifts_disagree(
                local_shift, global_shift,
                record, params['MAX_SHIFT_DISP_ALT']):
            corrected_shift = global_shift
        else:
            corrected_shift = local_shift

    corrected_shift = np.round(corrected_shift, 2)

    record.count_case(case)
    record.record_shift(corrected_shift, global_shift,
                        last_heads, local_shift, case)
    return corrected_shift

def predict_search_extent(obj1_extent, shift, params, grid_size):
    """ Predicts search extent/region for the object in image2 given
    the image shift. """
    shifted_center = obj1_extent['obj_center'] + shift
    search_radius_r = params['SEARCH_MARGIN'] / grid_size[1]
    search_radius_c = params['SEARCH_MARGIN'] / grid_size[2]
    x1 = shifted_center[0] - search_radius_r
    x2 = shifted_center[0] + search_radius_r + 1
    y1 = shifted_center[1] - search_radius_c
    y2 = shifted_center[1] + search_radius_c + 1
    x1 = np.int32(x1)
    x2 = np.int32(x2)
    y1 = np.int32(y1)
    y2 = np.int32(y2)
    return {'x1': x1, 'x2': x2, 'y1': y1, 'y2': y2,
            'center_pred': shifted_center, 'valid': True}

def check_search_box(search_box, img_dims):
    """ Checks if search_box is within the boundaries of the frame. Clips to
    edges of frame if out of bounds. Marks as invalid if too small. """
    if search_box['x1'] < 0:
        search_box['x1'] = 0
    if search_box['y1'] < 0:
        search_box['y1'] = 0
    if search_box['x2'] > img_dims[0]:
        search_box['x2'] = img_dims[0]
    if search_box['y2'] > img_dims[1]:
        search_box['y2'] = img_dims[1]
    if (
            (search_box['x2'] - search_box['x1'] < 5)
            or (search_box['y2'] - search_box['y1'] < 5)):
        search_box['valid'] = False
    return search_box

def get_disparity(obj_found, image2, search_box, obj1_extent):
    """ Computes disparities for objects in obj_found. """
    dist_pred = np.empty(0)
    change = np.empty(0)
    for target_obj in obj_found:
        target_extent = get_obj_extent(image2, target_obj)
        euc_dist = euclidean_dist(target_extent['obj_center'],
                                  search_box['center_pred'])
        dist_pred = np.append(dist_pred, euc_dist)
        size_changed = get_sizeChange(target_extent['obj_area'],
                                      obj1_extent['obj_area'])
        change = np.append(change, size_changed)
    # Note that merger of systems may create a sudden size change
    # that exaggerates cost function.
    disparity = dist_pred + change
    return disparity

def get_disparity_all(obj_found, image2, search_box, obj1_extent):
    """ Returns disparities of all objects found within the search box. """
    if np.max(obj_found) <= 0:
        disparity = np.array([LARGE_NUM])
    else:
        obj_found = obj_found[obj_found > 0]
        disparity = get_disparity(obj_found, image2,
                                  search_box, obj1_extent)
    return disparity

def save_obj_match(obj_id1, obj_found, disparity, obj_match, params):
    """ Saves disparity values in obj_match matrix. If disparity is greater
    than MAX_DISPARITY, saves a large number. """
    disparity[disparity > params['MAX_DISPARITY']] = LARGE_NUM
    if np.max(obj_found) > 0:
        obj_found = obj_found[obj_found > 0]
        obj_found = obj_found - 1
        obj_id1 = obj_id1 - 1
        obj_match[obj_id1, obj_found] = disparity
    return obj_match

def locate_all_objects(
        data_dic, global_shift, current_objects, record, params):
    """ Matches all the objects in image1 to objects in image2. This is the
    main function called on a pair of images. """
    nobj1 = np.max(data_dic['frame'])
    nobj2 = np.max(data_dic['frame_new'])

    if (nobj2 == 0) or (nobj1 == 0):
        print('No echoes to track!')
        return

    obj_match = np.full(
        (nobj1, np.max((nobj1, nobj2))), LARGE_NUM, dtype='f')
    u_shift = []
    v_shift = []

    for obj_id1 in np.arange(nobj1) + 1:
        obj1_extent = get_obj_extent(data_dic['frame'], obj_id1)

        shift = get_ambient_flow(
            obj1_extent, data_dic['refl'], data_dic['refl_new'], params,
            record.grid_size)

        if shift is None:
            record.count_case(5)
            shift = global_shift

        shift = correct_shift(
            shift, current_objects, obj_id1, global_shift, record, params)

        shift_meters = shift * record.grid_size[1:]
        [v, u] = shift_meters/record.interval.seconds
        u_shift.append(u)
        v_shift.append(v)

        search_box = predict_search_extent(
            obj1_extent, shift, params, record.grid_size)
        search_box = check_search_box(search_box, data_dic['frame_new'].shape)
        objs_found = find_objects(search_box, data_dic['frame_new'])
        disparity = get_disparity_all(
            objs_found, data_dic['frame_new'], search_box, obj1_extent)
        obj_match = save_obj_match(
            obj_id1, objs_found, disparity, obj_match, params)

    return obj_match, u_shift, v_shift

def match_pairs(obj_match, params):
    """ Matches objects into pairs given a disparity matrix and removes
    bad matches. Bad matches have a disparity greater than the maximum
    threshold. """

    # Create a list of sets, where the i-th set will store the objects
    # from image1 that have merged with objects in image2
    # Maybe faster to use a 2D array?
    obj_merge = np.zeros(obj_match.shape, dtype=bool)

    # Determine optimal pairs
    pairs = optimize.linear_sum_assignment(obj_match)

    for id1 in pairs[0]:
        if obj_match[id1, pairs[1][id1]] > params['MAX_DISPARITY']:
            # Set to -1 if object has died (or merged)
            pairs[1][id1] = -1
            # Find the closest object in image2 to object with id1
            id2 = np.argmin(obj_match[id1])
            # If this object was in the search radius of object id1,
            # add object id1 to obj_merge[id2].
            if obj_match[id1, id2] < LARGE_NUM:
                obj_merge[id1, id2] = True

    pairs = pairs[1] + 1  # ids in current_objects are 1-indexed
    return pairs, obj_merge

def get_pairs(data_dic, global_shift, current_objects, record, params):
    """ Given two images, this function identifies the matching objects and
    pairs them appropriately. See disparity function. """
    nobj1 = np.max(data_dic['frame'])
    nobj2 = np.max(data_dic['frame_new'])

    if nobj1 == 0:
        print('No echoes found in the first scan.')
        return
    elif nobj2 == 0:
        zero_pairs = np.zeros(nobj1)
        zero_obj_merge = np.zeros(
            (nobj1, np.max((nobj1, nobj2))), dtype=bool)
        return zero_pairs, zero_obj_merge, [np.nan] * nobj1, [np.nan] * nobj1

    obj_match, u_shift, v_shift = locate_all_objects(
        data_dic, global_shift, current_objects, record, params)

    pairs, obj_merge = match_pairs(obj_match, params)

    return pairs, obj_merge, u_shift, v_shift