#!/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.

import unittest
import copy
import numpy
import scipy.linalg

from pyscf import gto, lib
from pyscf import scf, dft
from pyscf.scf import addons

mol = gto.Mole()
mol.verbose = 7
mol.output = '/dev/null'
mol.atom = [
    ["O" , (0. , 0.     , 0.)],
    [1   , (0. , -0.757 , 0.587)],
    [1   , (0. , 0.757  , 0.587)] ]

mol.basis = {"H": '6-31g',
             "O": '6-31g',}
mol.build()
mf = scf.RHF(mol).run()

mol_dz = mol.copy()
mol_dz.basis = 'cc-pvdz'
mol_dz.cart = True
mol_dz.build(False, False)

mol1 = mol.copy()
mol1.spin = 2
mf_u = scf.UHF(mol1).run()

mol2 = mol.copy()
mol2.symmetry = True
mol2.build(0,0)
sym_mf = scf.RHF(mol2).run()

mol3 = mol1.copy()
mol3.symmetry = True
mol3.spin = 2
mol3.build(0,0)
sym_mf_u = scf.UHF(mol3).run()

def tearDownModule():
    global mol, mf, mol_dz, mol1, mf_u, mol2, sym_mf, mol3, sym_mf_u
    mol.stdout.close()
    del mol, mf, mol_dz, mol1, mf_u, mol2, sym_mf, mol3, sym_mf_u


class KnownValues(unittest.TestCase):
    def test_project_mo_nr2nr(self):
        nao = mol.nao_nr()
        c = numpy.random.random((nao,nao))
        c1 = addons.project_mo_nr2nr(mol, c, mol)
        self.assertAlmostEqual(abs(c-c1).max(), 0, 12)

        numpy.random.seed(15)
        nao = mol.nao_nr()
        mo1 = numpy.random.random((nao,nao))
        mo2 = addons.project_mo_nr2nr(mol, [mo1,mo1], mol_dz)
        self.assertAlmostEqual(abs(mo2[0]).sum(), 83.436359425591888, 11)
        self.assertAlmostEqual(abs(mo2[1]).sum(), 83.436359425591888, 11)

    def test_project_mo_r2r(self):
        nao = mol.nao_2c()
        c = numpy.random.random((nao*2,nao*2))
        c = c + numpy.sin(c)*1j
        c1 = addons.project_mo_r2r(mol, c, mol)
        self.assertAlmostEqual(abs(c-c1).max(), 0, 12)

        numpy.random.seed(15)
        n2c = mol.nao_2c()
        n4c = n2c * 2
        mo1 = numpy.random.random((n4c,n4c)) + numpy.random.random((n4c,n4c))*1j
        mo2 = addons.project_mo_r2r(mol, [mo1,mo1], mol_dz)
        self.assertAlmostEqual(abs(mo2[0]).sum(), 2159.3715489514038, 11)
        self.assertAlmostEqual(abs(mo2[1]).sum(), 2159.3715489514038, 11)

    def test_project_mo_nr2r(self):
        numpy.random.seed(15)
        nao = mol.nao_nr()
        mo1 = numpy.random.random((nao,nao))
        mo2 = addons.project_mo_nr2r(mol, [mo1,mo1], mol_dz)
        self.assertAlmostEqual(abs(mo2[0]).sum(), 172.66468850263556, 11)
        self.assertAlmostEqual(abs(mo2[1]).sum(), 172.66468850263556, 11)

        mo2 = addons.project_mo_nr2r(mol, mo1, mol_dz)
        self.assertAlmostEqual(abs(mo2).sum(), 172.66468850263556, 11)

    def test_project_dm_nr2nr(self):
        nao = mol.nao_nr()
        dm = numpy.random.random((nao,nao))