# Copyright 2014-2020 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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from pyscf import scf
from pyscf.tdscf import rhf
from pyscf.tdscf import uhf
from pyscf.tdscf import ghf
from pyscf.tdscf import dhf
from pyscf.tdscf.rhf import TDRHF
from pyscf.tdscf.uhf import TDUHF
from pyscf.tdscf.ghf import TDGHF

try:
    from pyscf.dft import KohnShamDFT
    from pyscf.tdscf import rks
    from pyscf.tdscf import uks
    from pyscf.tdscf import gks
    from pyscf.tdscf import dks
    from pyscf.tdscf.rks import TDRKS
    from pyscf.tdscf.uks import TDUKS
    from pyscf.tdscf.gks import TDGKS
except (ImportError, IOError):
    pass


def TDHF(mf, frozen=None):
    if isinstance(mf, scf.hf.KohnShamDFT):
        raise RuntimeError('TDHF does not support DFT object %s' % mf)
    mf = mf.remove_soscf()
    if isinstance(mf, scf.rohf.ROHF):
        # Is it correct to call TDUHF for ROHF?
        mf = mf.to_uhf()
    return mf.TDHF(frozen)

def TDA(mf, frozen=None):
    mf = mf.remove_soscf()
    if isinstance(mf, scf.rohf.ROHF):
        if isinstance(mf, KohnShamDFT):
            mf = mf.to_uks()
        else:
            mf = mf.to_uhf()
    return mf.TDA(frozen)

def TDDFT(mf, frozen=None):
    if isinstance(mf, KohnShamDFT):
        mf = mf.remove_soscf()
        if isinstance(mf, scf.rohf.ROHF):
            mf = mf.to_uks()
        return mf.TDDFT(frozen)
    else:
        return TDHF(mf, frozen)

TD = TDDFT


def RPA(mf, frozen=None):
    return TDDFT(mf, frozen)

def dRPA(mf, frozen=None):
    mf = mf.remove_soscf()
    if isinstance(mf, scf.rohf.ROHF):
        if isinstance(mf, KohnShamDFT):
            mf = mf.to_uks()
        else:
            mf = mf.to_uhf()
    return mf.dRPA(frozen)

def dTDA(mf, frozen=None):
    mf = mf.remove_soscf()
    if isinstance(mf, scf.rohf.ROHF):
        if isinstance(mf, KohnShamDFT):
            mf = mf.to_uks()
        else:
            mf = mf.to_uhf()
    return mf.dTDA(frozen)
