#!/usr/bin/env python
# coding: utf-8

# ## Initial imports

import os
import sys

import matplotlib.pyplot as plt
import numpy as np

from scm.conformers import ConformersJob
from scm.plams import Settings, view


# ## Initial structure

from scm.base import ChemicalSystem

molecule = ChemicalSystem.from_smiles("OC(CC1c2ccccc2Sc2ccccc21)CN1CCCC1")
view(molecule, width=300, height=300, direction="along_pca3", picture_path="picture1.png")


# ## Generate conformers with RDKit and UFF
# The fastest way to generate conformers is to use RDKit with the UFF force field.
#
# Below we specify to generate 16 initial conformers. The final number of conformers may be smaller, as the geometry optimization may cause several structures to enter the same minimum.

# ### Conformer generation settings

s = Settings()
s.input.ams.Task = "Generate"  # default
s.input.ams.Generator.Method = "RDKit"  # default
s.input.ams.Generator.RDKit.InitialNConformers = 16  # optional, non-default
s.input.ForceField.Type = "UFF"  # default


# ### Conformer generation input file

print(ConformersJob(settings=s).get_input())


# ### Run conformer generation

generate_job = ConformersJob(name="generate", molecule=molecule, settings=s)
generate_job.run()
# ## Conformer generation results

# ### Some helper functions

import pandas as pd


def get_df(job: ConformersJob, temperature=298, unit="kcal/mol") -> pd.DataFrame:
    relative_energies = job.results.get_relative_energies(unit)
    pop = job.results.get_boltzmann_distribution(temperature)
    df = pd.DataFrame(
        {
            "Conformer Id": np.arange(1, len(relative_energies) + 1),
            f"Rel. E ({unit})": relative_energies,
            f"Population at {temperature} K": pop,
        }
    )
    return df


# ### Actual results
#
# Below we see that the **conformer generation gave 15 distinct conformers**, where the highest-energy conformer is 18 kcal/mol higher in energy than the lowest energy conformer.
#
# You can also see the **relative populations** of these conformers at the specified temperature. The populations are calculated from the **Boltzmann distribution** and the relative energies.

unit = "kcal/mol"
temperature = 298


df = get_df(generate_job, temperature, unit)
print(df)


from scm.conformers.plams.plot import plot_conformers

axes = plot_conformers(
    generate_job,
    4,
    temperature=temperature,
    unit=unit,
    lowest=True,
    width=300,
    height=300,
    direction="along_pca3",
)
axes
import matplotlib.pyplot as plt

plt.gcf().savefig("picture2.png")


# ## Re-optimize conformers with GFNFF
#
# The UFF force field is not very accurate for geometries and energies. From an initial conformer set you can reoptimize it with a better level of theory.
#
# The **Optimize** task performs **GeometryOptimization** jobs on each conformer in a set.
#
# Below, the most stable conformers (within 8 kcal/mol of the most stable conformer) at the UFF level of theory are re-optimized with GFNFF, which gives more accurate geometries.

s = Settings()
s.input.ams.Task = "Optimize"
# must be absolute path
s.input.ams.InputConformersSet = os.path.abspath(generate_job.results.rkfpath())
# only conformers within 8 kcal/mol at the PREVIOUS level of theory
s.input.ams.InputMaxEnergy = 8.0
s.input.GFNFF  # or choose a different engine if you don't have a GFNFF license

reoptimize_job = ConformersJob(settings=s, name="reoptimize")
print(reoptimize_job.get_input())


reoptimize_job.run()
df = get_df(reoptimize_job, temperature=temperature, unit=unit)
print(df)


axes = plot_conformers(reoptimize_job, 4, temperature=temperature, unit=unit, lowest=True)
axes
import matplotlib.pyplot as plt

plt.gcf().savefig("picture3.png")


# ## Score conformers with DFTB
#
# If you have many conformers or a very large molecule, it can be computationally expensive to do the conformer generation or reoptimization and a high level of theory.
#
# The **Score** task runs **SinglePoint** jobs on the conformers in a set. This lets you use a more computationally expensive method. Here, we choose DFTB, although normally you may choose some DFT method.

s = Settings()
s.input.ams.Task = "Score"
# must be absolute path
s.input.ams.InputConformersSet = os.path.abspath(reoptimize_job.results.rkfpath())
# only conformers within 4 kcal/mol at the PREVIOUS level of theory
s.input.ams.InputMaxEnergy = 4.0
s.input.DFTB.Model = "GFN1-xTB"  # or choose a different engine if you don't have a DFTB license
# s.input.adf.XC.GGA = 'PBE'                       # to use ADF PBE
# s.input.adf.XC.DISPERSION = 'GRIMME3 BJDAMP'     # to use ADF PBE with Grimme D3(BJ) dispersion

score_job = ConformersJob(settings=s, name="score")
score_job.run()
df = get_df(score_job, temperature=temperature, unit=unit)
print(df)


axes = plot_conformers(score_job, 4, temperature=temperature, unit=unit, lowest=True)
axes
import matplotlib.pyplot as plt

plt.gcf().savefig("picture4.png")


# Here, you see that from the conformers in the set, **DFTB predicts a different lowest-energy conformer than GFNFF** (compare to previous figure).

# ## Filter a conformer set
#
# In practice, you may have generated thousands of conformers for a particular structure. Many of those conformers may be so high in energy that their Boltzmann weights are very small.
#
# The **Filter** task only filters the conformers, it does not perform any additional calculations. It can be used to reduce a conformer set so that it is more convenient to work with.
#
# Below, we filter the conformers set to only the conformers within 2 kcal/mol of the minimum.

s = Settings()
s.input.ams.Task = "Filter"
s.input.ams.InputConformersSet = os.path.abspath(score_job.results.rkfpath())
s.input.ams.InputMaxEnergy = 2.0

filter_job = ConformersJob(settings=s, name="filter")
filter_job.run()
df = get_df(filter_job, temperature=temperature, unit=unit)
print(df)


axes = plot_conformers(filter_job, 4, temperature=temperature, unit=unit, lowest=True)
axes
import matplotlib.pyplot as plt

plt.gcf().savefig("picture5.png")


# The structures and energies are identical to before. However, the relative populations changed slightly as there are now fewer conformers in the set.

# ## More about conformers
#
# * Try **CREST** instead of RDKit to generate the initial conformer set
#
# * The **Expand** task can be used to expand a set of conformers.