"""Basic example script.

This script is a modified version from the ground state cas calculation
in scine_autocas.main_functions
"""
# -*- coding: utf-8 -*-
__copyright__ = """This file is part of SCINE AutoCAS.
This code is licensed under the 3-clause BSD license.
Copyright ETH Zurich, Laboratory for Physical Chemistry, Reiher Group.
See LICENSE.txt for details
"""

import os

from scine_autocas import Autocas
from scine_autocas.autocas_utils.molecule import Molecule
from scine_autocas.interfaces.molcas import Molcas
from scine_autocas.main_functions import MainFunctions
from scine_autocas.plots.entanglement_plot import EntanglementPlot
from scine_autocas.plots.threshold_plot import ThresholdPlot


def standard_autocas_procedure(path: str, xyz_file: str):
    """Do ground state cas"""
    # looks better in output (escape \)
    print("                                  _            _____             _____                     ")
    print("                                 | |          / ____|    /\\     / ____|                    ")
    print("                     __ _  _   _ | |_   ___  | |        /  \\   | (___                      ")
    print("                    / _` || | | || __| / _ \\ | |       / /\\ \\   \\___ \\                     ")
    print("                   | (_| || |_| || |_ | (_) || |____  / ____ \\  ____) |                    ")
    print("                    \\__,_| \\__,_| \\__| \\___/  \\_____|/_/    \\_\\|_____/                     ")
    print("                                                                                           ")

    print("*******************************************************************************************")
    print("*                                                                                         *")
    print("*                                      Settings                                           *")
    print("*                                                                                         *")
    print("*******************************************************************************************")
    # create a molecule
    molecule = Molecule(xyz_file)

    # initialize autoCAS and Molcas interface
    autocas = Autocas(molecule)
    molcas = Molcas(molecules=[molecule])

    # setup interface
    molcas.project_name = "autocas"
    molcas.settings.work_dir = path 
    molcas.environment.molcas_scratch_dir = path + "/molcas_scratch"
    molcas.settings.xyz_file = xyz_file
    molcas.settings.orbital_localisation = True
    molcas.settings.localisation_space = "ALL"
    molcas.settings.basis_set = "STO-3G"
    print('Molcas Settings:')
    print('localisation',molcas.settings.orbital_localisation)
    print('basis', molcas.settings.basis_set)
    
    # make initial HF calculation
    print("")
    print("*******************************************************************************************")
    print("*                                                                                         *")
    print("*                                  Initial HF Calculation                                 *")
    print("*                                                                                         *")
    print("*******************************************************************************************")
    molcas.calculate()

    # make DMRG calculation with large active space
    molcas.settings.method = "DMRGCI"

    # manually set dmrg sweeps and bond dmrg_bond_dimension to low number
    molcas.settings.dmrg_bond_dimension = 250
    molcas.settings.dmrg_sweeps = 5

    # make initial active space and evaluate initial DMRG calculation
    occ_initial, index_initial = autocas.make_initial_active_space()
    print("")
    print("*******************************************************************************************")
    print("*                                                                                         *")
    print("*                                    DMRG Calculation                                     *")
    print("*                                                                                         *")
    print("*******************************************************************************************")

    print("Settings for autoCAS DMRG calculation:")
    print("method:", molcas.settings.method)
    print("dmrg_sweeps:", molcas.settings.dmrg_sweeps)
    print("dmrg_bond_dimension: ", molcas.settings.dmrg_bond_dimension)
    print("")
    
    # do cas calculation
    cas_results = molcas.calculate(occ_initial, index_initial)

    print(f"Initial active space CAS(e, o): ({sum(occ_initial)}, {len(occ_initial)})")
    print(f"Orbital indices:     {index_initial}")
    print(f"Orbital occupations: {occ_initial}")
    print("")
        
    energy = cas_results[0]
    s1_entropy = cas_results[1]
    s2_entropy = cas_results[2]
    mut_inf = cas_results[3]

    # plot entanglement diagram
    print("Creating entanglement diagram.")
    entanglement_plot = EntanglementPlot()
    plt = entanglement_plot.plot(s1_entropy, mut_inf)  # type: ignore
    plt.savefig(molcas.settings.work_dir + "/entanglement_diagram.pdf")  # type: ignore
    
    print("Creating threshold diagram.")
    threshold_plot = ThresholdPlot()
    plt = threshold_plot.plot(s1_entropy)
    plt.savefig(molcas.settings.work_dir + "/threshold_diagram.pdf")

    # make active space based on single orbital entropies
    cas_occ, cas_index = autocas.get_active_space(
        occ_initial, s1_entropy   # type: ignore
    )
    print('only active space search is activated, no final CAS energy is calculated')
    final_energy=1.0
    return cas_occ, cas_index, final_energy


if __name__ == "__main__":
    path_to_this_file = os.path.dirname(os.path.abspath(__file__))
    xyz = path_to_this_file + "/n2.xyz"
    occupation, orbitals, energy = standard_autocas_procedure(path_to_this_file, xyz)