from src import scf
from src.utils import system
import numpy as np
from mindspore import Tensor
from mindspore.ops import composite as C
from mindspore import nn

grad_all_with_sens = C.GradOperation(get_all=True, sens_param=True)
random_normal = C.normal


class GradNet(nn.Cell):
    def __init__(self, net):
        super(GradNet, self).__init__()
        self.net = net

    def construct(self, x, sens):
        return grad_all_with_sens(self.net)(x, sens)


def create_o2_hf(bond_length):
    molecule = [
        system.Atom('O', (0, 0, 0)),
        system.Atom('O', (0, 0, bond_length))
    ]
    spin = 2
    oxygen_atomic_number = 8
    nelectrons = [oxygen_atomic_number + spin, oxygen_atomic_number - spin]
    hf = scf.Scf(molecule=molecule, nelectrons=nelectrons, restricted=False)
    return hf


def test_ms_eval_mos():
    xs = np.random.randn(100, 3)
    hf = create_o2_hf(1.4)
    hf.run()

    mo_vals = [mo_val for mo_val in hf.eval_mos(xs, deriv=True)]
    ms_res = hf.ms_eval_mos(Tensor(xs), deriv=True)
    gnet = GradNet(hf.net)

    ms_dmo_vals = gnet(Tensor(xs), Tensor(np.concatenate(mo_vals, axis=-1)[0]))
    ms_res = np.split(ms_res.asnumpy(), 2, axis=-1)
    for np_val, ms_val in zip(mo_vals, ms_res):
        np.testing.assert_allclose(np_val[0], ms_val)
    np_deriv = sum(
        np.sum(np_val[1:] * np.expand_dims(np_val[0], 0), axis=-1)
        for np_val in mo_vals)
    np_deriv = np.transpose(np_deriv)
    np.testing.assert_allclose(np_deriv, ms_dmo_vals[0].asnumpy(), atol=1e-6, rtol=1e-6)


def test_ms_eval_mos_deriv():
    hf = create_o2_hf(1.4)
    hf.run()
    xs = random_normal((100, 3), Tensor(0), Tensor(1))
    ms_mos_derivs = hf.ms_eval_mos(xs, deriv=True)
    ms_mos = hf.ms_eval_mos(xs, deriv=False)
    np.testing.assert_allclose(ms_mos.asnumpy(), ms_mos_derivs.asnumpy())


def test_ms_eval_hf():
    molecule = [system.Atom('O', (0, 0, 0))]
    nelectrons = (5, 3)
    hf = scf.Scf(molecule=molecule,
                 nelectrons=nelectrons,
                 restricted=False)
    hf.run()

    batch = 100
    xs = [np.random.randn(batch, 3) for _ in range(sum(nelectrons))]
    mos = []
    for i, x in enumerate(xs):
        ispin = 0 if i < nelectrons[0] else 1
        orbitals = hf.eval_mos(x)[ispin]
        mos.append(orbitals[:, :nelectrons[ispin]])
    np_mos = (np.stack(mos[:nelectrons[0]], axis=1),
              np.stack(mos[nelectrons[0]:], axis=1))

    ms_xs = Tensor(np.stack(xs, axis=1))
    ms_mos = hf.ms_eval_hf(ms_xs, deriv=True)

    for i, (np_mos_mat, ms_mos_mat) in enumerate(zip(np_mos, ms_mos)):
        ms_mos_mat_np = ms_mos_mat.asnumpy()
        assert np_mos_mat.shape == ms_mos_mat_np.shape
        assert np_mos_mat.shape == (batch, nelectrons[i], nelectrons[i])
        np.testing.assert_allclose(np_mos_mat, ms_mos_mat_np)


def test_ms_eval_slog_wavefuncs():
    molecule = [system.Atom('O', (0, 0, 0))]
    nelectrons = (5, 3)
    total_electrons = sum(nelectrons)
    num_spatial_dim = 3
    hf = scf.Scf(molecule=molecule,
                 nelectrons=nelectrons,
                 restricted=False)
    hf.run()

    batch = 100
    rng = np.random.RandomState(1)
    flat_positions_np = rng.randn(batch,
                                  total_electrons * num_spatial_dim)

    flat_positions_ms = Tensor(flat_positions_np)
    for method in [hf.ms_eval_slog_slater_determinant,
                   hf.ms_eval_slog_hartree_product]:
        slog_wavefunc, signs = method(flat_positions_ms)
        slog_wavefunc_ = slog_wavefunc.asnumpy()
        signs_ = signs.asnumpy()
        assert slog_wavefunc_.shape == (batch, 1)
        assert signs_.shape == (batch, 1)
    hartree_product = hf.ms_eval_hartree_product(flat_positions_ms)
    hartree_product_ = hartree_product.asnumpy()
    np.testing.assert_almost_equal(hartree_product_,
                                   np.exp(slog_wavefunc_) * signs_)


test_ms_eval_mos()
test_ms_eval_mos_deriv()
test_ms_eval_hf()
test_ms_eval_slog_wavefuncs()