#!/usr/bin/env python
# Copyright 2020-2023 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: Xing Zhang <zhangxing.nju@gmail.com>
#

import unittest
import numpy as np

from pyscf.pbc import gto as pbcgto
from pyscf.pbc import scf as pscf
from pyscf.pbc.dft import krks,kuks,multigrid

def make_primitive_cell(mesh):
    cell = pbcgto.Cell()
    cell.unit = 'A'
    cell.atom = 'Si 0.,  0.,  0.; Si 1.3467560987,  1.3467560987,  1.3467560987'
    cell.a = '''0.            2.6935121974    2.6935121974
                2.6935121974  0.              2.6935121974
                2.6935121974  2.6935121974    0.    '''

    cell.basis = 'gth-szv'
    cell.pseudo = 'gth-pade'
    cell.mesh = mesh
    cell.spin = 0
    cell.verbose = 0
    cell.output = '/dev/null'
    cell.space_group_symmetry = True
    cell.build()
    return cell

def setUpModule():
    global cell, He, nk
    cell = make_primitive_cell([16]*3)
    nk = [1,2,2]

    L = 2.
    He = pbcgto.Cell()
    He.verbose = 0
    He.a = np.eye(3)*L
    He.atom =[['He' , ( L/2+0., L/2+0., L/2+0.)],]
    He.basis = {'He': [[0, (4.0, 1.0)], [0, (1.0, 1.0)]]}
    He.space_group_symmetry = True
    He.build()

def tearDownModule():
    global cell, He, nk
    del cell, He, nk

class KnownValues(unittest.TestCase):
    def test_krks_gamma_center(self):
        kpts0 = cell.make_kpts(nk, with_gamma_point=True)
        kmf0 = krks.KRKS(cell, kpts=kpts0)
        kmf0.xc = 'lda'
        kmf0.kernel()
        rho0 = kmf0.get_rho()

        kpts = cell.make_kpts(nk, with_gamma_point=True,space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(cell, kpts=kpts)
        kmf.xc = 'lda'
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 7)
        rho = kmf.get_rho()
        error = np.amax(np.absolute(rho - rho0))
        self.assertAlmostEqual(error, 0., 7)

    def test_krks_monkhorst(self):
        kpts0 = cell.make_kpts(nk, with_gamma_point=False)
        kmf0 = krks.KRKS(cell, kpts=kpts0)
        kmf0.xc = 'lda'
        kmf0.kernel()

        kpts = cell.make_kpts(nk, with_gamma_point=False,space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(cell, kpts=kpts)
        kmf.xc = 'lda'
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 6)

    def test_kuks_gamma_center(self):
        kpts0 = cell.make_kpts(nk, with_gamma_point=True)
        kumf0 = kuks.KUKS(cell, kpts=kpts0)
        kumf0.xc = 'lda'
        kumf0.kernel()
        rho0 = kumf0.get_rho()

        kpts = cell.make_kpts(nk, with_gamma_point=True,space_group_symmetry=True,time_reversal_symmetry=True)
        kumf = pscf.KUKS(cell, kpts=kpts)
        kumf.xc = 'lda'
        kumf.kernel()
        rho = kumf.get_rho()
        self.assertAlmostEqual(kumf.e_tot, kumf0.e_tot, 7)
        error = np.amax(np.absolute(rho - rho0))
        self.assertAlmostEqual(error, 0., 7)

    def test_kuks_monkhorst(self):
        kpts0 = cell.make_kpts(nk, with_gamma_point=False)
        kumf0 = kuks.KUKS(cell, kpts=kpts0)
        kumf0.xc = 'lda'
        kumf0.kernel()

        kpts = cell.make_kpts(nk, with_gamma_point=False,space_group_symmetry=True,time_reversal_symmetry=True)
        kumf = pscf.KUKS(cell, kpts=kpts)
        kumf.xc = 'lda'
        kumf.kernel()
        self.assertAlmostEqual(kumf.e_tot, kumf0.e_tot, 6)

    def test_rsh(self):
        kpts0 = He.make_kpts(nk, with_gamma_point=False)
        kmf0 = krks.KRKS(He, kpts=kpts0)
        kmf0.xc = 'camb3lyp'
        kmf0.kernel()

        kpts = He.make_kpts(nk, with_gamma_point=False, space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(He, kpts=kpts)
        kmf.xc = 'camb3lyp'
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 9)

    def test_lda_df(self):
        kpts0 = He.make_kpts(nk, with_gamma_point=False)
        kmf0 = krks.KRKS(He, kpts=kpts0).density_fit()
        kmf0.xc = 'lda'
        kmf0.kernel()

        kpts = He.make_kpts(nk, with_gamma_point=False, space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(He, kpts=kpts).density_fit()
        kmf.xc = 'lda'
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 9)

    def test_gga_df(self):
        kpts0 = He.make_kpts(nk, with_gamma_point=False)
        kmf0 = krks.KRKS(He, kpts=kpts0).density_fit()
        kmf0.xc = 'pbe'
        kmf0.kernel()

        kpts = He.make_kpts(nk, with_gamma_point=False, space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(He, kpts=kpts).density_fit()
        kmf.xc = 'pbe'
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 9)

    def test_gga_df_newton(self):
        kpts0 = He.make_kpts(nk, with_gamma_point=False)
        kmf0 = krks.KRKS(He, kpts=kpts0).density_fit()
        kmf0.xc = 'pbe'
        kmf0 = kmf0.newton()
        kmf0.kernel()

        kpts = He.make_kpts(nk, with_gamma_point=False, space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(He, kpts=kpts).density_fit()
        kmf.xc = 'pbe'
        kmf = kmf.newton()
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 9)

    def test_rsh_df(self):
        kpts0 = He.make_kpts(nk, with_gamma_point=False)
        kmf0 = krks.KRKS(He, kpts=kpts0).density_fit()
        kmf0.xc = 'camb3lyp'
        kmf0.kernel()

        kpts = He.make_kpts(nk, with_gamma_point=False, space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(He, kpts=kpts).density_fit()
        kmf.xc = 'camb3lyp'
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 8)

    def test_rsh_mdf(self):
        kpts0 = He.make_kpts(nk, with_gamma_point=False)
        kmf0 = krks.KRKS(He, kpts=kpts0).mix_density_fit()
        kmf0.xc = 'camb3lyp'
        kmf0.kernel()

        kpts = He.make_kpts(nk, with_gamma_point=False, space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(He, kpts=kpts).mix_density_fit()
        kmf.xc = 'camb3lyp'
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 8)

    def test_multigrid(self):
        kmf0 = krks.KRKS(cell, kpts=cell.make_kpts(nk))
        kmf0.xc = 'lda'
        kmf0 = multigrid.multigrid(kmf0)
        kmf0.kernel()
        rho0 = kmf0.get_rho()

        kpts = cell.make_kpts(nk,space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(cell, kpts=kpts)
        kmf.xc = 'lda'
        kmf = multigrid.multigrid(kmf)
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 7)
        rho = kmf.get_rho()
        error = np.amax(np.absolute(rho - rho0))
        self.assertAlmostEqual(error, 0., 7)

        kmf.with_df = multigrid.MultiGridFFTDF(cell, kpts)
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 7)
        rho = kmf.get_rho()
        error = np.amax(np.absolute(rho - rho0))
        self.assertAlmostEqual(error, 0., 7)

    def test_multigrid_kuks(self):
        kmf0 = pscf.KUKS(cell, kpts=cell.make_kpts(nk))
        kmf0.xc = 'lda'
        kmf0 = multigrid.multigrid(kmf0)
        kmf0.kernel()
        rho0 = kmf0.get_rho()

        kpts = cell.make_kpts(nk,space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KUKS(cell, kpts=kpts)
        kmf.xc = 'lda'
        kmf = multigrid.multigrid(kmf)
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 7)
        rho = kmf.get_rho()
        error = np.amax(np.absolute(rho - rho0))
        self.assertAlmostEqual(error, 0., 7)

        kmf.with_df = multigrid.MultiGridFFTDF(cell, kpts)
        kmf.kernel()
        self.assertAlmostEqual(kmf.e_tot, kmf0.e_tot, 7)
        rho = kmf.get_rho()
        error = np.amax(np.absolute(rho - rho0))
        self.assertAlmostEqual(error, 0., 7)

    def test_to_uhf(self):
        kpts = cell.make_kpts(nk,space_group_symmetry=True,time_reversal_symmetry=True)
        kmf = pscf.KRKS(cell, kpts=kpts)
        kmf.xc = 'lda'
        kmf.kernel()
        dm = kmf.make_rdm1()
        dm = np.asarray([dm,dm]) / 2.

        kumf = kmf.to_uhf()
        kumf.max_cycle = 1
        kumf.kernel(dm)
        self.assertAlmostEqual(kmf.e_tot, kumf.e_tot, 8)

    def test_to_rhf(self):
        kpts = cell.make_kpts(nk,space_group_symmetry=True,time_reversal_symmetry=True)
        kumf = pscf.KUKS(cell, kpts=kpts)
        kumf.xc = 'lda'
        kumf.kernel()
        dm = kumf.make_rdm1()

        kmf = kumf.to_rhf()
        kmf.max_cycle = 1
        kmf.kernel(dm[0]+dm[1])
        self.assertAlmostEqual(kmf.e_tot, kumf.e_tot, 8)

    def test_convert_from(self):
        kpts = cell.make_kpts(nk,space_group_symmetry=True,time_reversal_symmetry=True)
        kumf = pscf.KUKS(cell, kpts=kpts)
        kumf.xc = 'lda'
        kumf.kernel()
        dm = kumf.make_rdm1()

        kmf = pscf.KRKS(cell, kpts=kpts)
        kmf = kmf.convert_from_(kumf)
        kmf.max_cycle = 1
        kmf.kernel(dm[0]+dm[1])
        self.assertAlmostEqual(kmf.e_tot, kumf.e_tot, 8)

        dm = kmf.make_rdm1()
        kumf = kumf.convert_from_(kmf)
        kumf.max_cycle = 1
        kumf.kernel(np.asarray([dm,dm]) / 2.)
        self.assertAlmostEqual(kmf.e_tot, kumf.e_tot, 7)

if __name__ == '__main__':
    print("Full Tests for DFT with k-point symmetry")
    unittest.main()
