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

from functools import reduce
import unittest
import numpy
import scipy.linalg
from pyscf import lib
from pyscf import gto
from pyscf.x2c import sfx2c1e
from pyscf.x2c import sfx2c1e_grad

def _sqrt0(a):
    w, v = scipy.linalg.eigh(a)
    return numpy.dot(v*numpy.sqrt(w), v.conj().T)

def _invsqrt0(a):
    w, v = scipy.linalg.eigh(a)
    return numpy.dot(v/numpy.sqrt(w), v.conj().T)

def _sqrt1(a0, a1):
    '''Solving first order of x^2 = a'''
    w, v = scipy.linalg.eigh(a0)
    w = numpy.sqrt(w)
    a1 = reduce(numpy.dot, (v.conj().T, a1, v))
    x1 = a1 / (w[:,None] + w)
    x1 = reduce(numpy.dot, (v, x1, v.conj().T))
    return x1

def _invsqrt1(a0, a1):
    '''Solving first order of x^2 = a^{-1}'''
    w, v = scipy.linalg.eigh(a0)
    w = 1./numpy.sqrt(w)
    a1 = -reduce(numpy.dot, (v.conj().T, a1, v))
    x1 = numpy.einsum('i,ij,j->ij', w**2, a1, w**2) / (w[:,None] + w)
    x1 = reduce(numpy.dot, (v, x1, v.conj().T))
    return x1

def get_R(mol):
    s0 = mol.intor('int1e_ovlp')
    t0 = mol.intor('int1e_kin')
    s0sqrt = _sqrt0(s0)
    s0invsqrt = _invsqrt0(s0)