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

# ## Initial imports

from scm.plams import plot_molecule, from_smiles, Molecule, packmol, packmol_around, view
from ase.visualize.plot import plot_atoms
from ase.build import fcc111, bulk
import matplotlib.pyplot as plt


# ## Helper functions


def printsummary(mol, details=None):
    if details:
        density = details["density"]
    else:
        density = mol.get_density() * 1e-3
    s = f"{len(mol)} atoms, density = {density:.3f} g/cm^3"
    if mol.lattice:
        s += f", box = {mol.lattice[0][0]:.3f}, {mol.lattice[1][1]:.3f}, {mol.lattice[2][2]:.3f}"
    s += f", formula = {mol.get_formula()}"
    if details:
        s += f"\n#added molecules per species: {details['n_molecules']}, mole fractions: {details['mole_fractions']}"
    print(s)


# ## Liquid water (fluid with 1 component)
# First, create the gasphase molecule:

water = from_smiles("O")
view(water, width=200, height=200, picture_path="picture1.png")


print("pure liquid from approximate number of atoms and exact density (in g/cm^3), cubic box with auto-determined size")
out = packmol(water, n_atoms=194, density=1.0)
printsummary(out)
out.write("water-1.xyz")
view(out, width=300, height=300, padding=-2, picture_path="picture2.png")


print("pure liquid from approximate density (in g/cm^3) and an orthorhombic box")
out = packmol(water, density=1.0, box_bounds=[0.0, 0.0, 0.0, 8.0, 12.0, 14.0])
printsummary(out)
out.write("water-2.xyz")
view(out, width=300, height=300, padding=-2, picture_path="picture3.png")


print("pure liquid with explicit number of molecules and exact density")
out = packmol(water, n_molecules=64, density=1.0)
printsummary(out)
out.write("water-3.xyz")
view(out, width=300, height=300, padding=-2, picture_path="picture4.png")


print("pure liquid with explicit number of molecules and box")
out = packmol(water, n_molecules=64, box_bounds=[0.0, 0.0, 0.0, 12.0, 13.0, 14.0])
printsummary(out)
out.write("water-4.xyz")
view(out, width=300, height=300, padding=-2, picture_path="picture5.png")


print("water-5.xyz: pure liquid in non-orthorhombic box (requires AMS2025 or later)")
print("NOTE: Non-orthorhombic boxes may yield inaccurate results, always carefully check the output")
# You can pack inside any lattice using the packmol_around function
box = Molecule()
box.lattice = [[10.0, 2.0, -1.0], [-5.0, 8.0, 0.0], [0.0, -2.0, 11.0]]
out = packmol_around(box, molecules=[water], n_molecules=[32])
out.write("water-5.xyz")
view(out, width=300, height=300, padding=-2, picture_path="picture6.png")


print("Experimental feature (AMS2025): guess density for pure liquid")
print("Note: This density is meant to be equilibrated with NPT MD. It can be very inaccurate!")
out = packmol(water, n_atoms=100)
print(f"Guessed density: {out.get_density():.2f} kg/m^3")
view(out, width=200, height=200, padding=-2, picture_path="picture7.png")


# ## Water-acetonitrile mixture (fluid with 2 or more components)
# Let's also create a single acetonitrile molecule:

acetonitrile = from_smiles("CC#N")
view(acetonitrile, width=200, height=200, picture_path="picture8.png")


# Set the desired mole fractions and density. Here, the density is calculated as the weighted average of water (1.0 g/cm^3) and acetonitrile (0.76 g/cm^3) densities, but you could use any other density.

# MIXTURES
x_water = 0.666  # mole fraction
x_acetonitrile = 1 - x_water  # mole fraction
# weighted average of pure component densities
density = (x_water * 1.0 + x_acetonitrile * 0.76) / (x_water + x_acetonitrile)

print("MIXTURES")
print(f"x_water = {x_water:.3f}")
print(f"x_acetonitrile = {x_acetonitrile:.3f}")
print(f"target density = {density:.3f} g/cm^3")


# By setting ``return_details=True``, you can get information about the mole fractions of the returned system. They may not exactly match the mole fractions you put in.

print(
    "2-1 water-acetonitrile from approximate number of atoms and exact density (in g/cm^3), "
    "cubic box with auto-determined size"
)
out, details = packmol(
    molecules=[water, acetonitrile],
    mole_fractions=[x_water, x_acetonitrile],
    n_atoms=200,
    density=density,
    return_details=True,
)
printsummary(out, details)
out.write("water-acetonitrile-1.xyz")
view(out, width=300, height=300, padding=-2, picture_path="picture9.png")


# The ``details`` is a dictionary as follows:

for k, v in details.items():
    print(f"{k}: {v}")


print("2-1 water-acetonitrile from approximate density (in g/cm^3) and box bounds")
out, details = packmol(
    molecules=[water, acetonitrile],
    mole_fractions=[x_water, x_acetonitrile],
    box_bounds=[0, 0, 0, 13.2, 13.2, 13.2],
    density=density,
    return_details=True,
)
printsummary(out, details)
out.write("water-acetonitrile-2.xyz")
view(out, width=300, height=300, padding=-2, picture_path="picture10.png")


print("2-1 water-acetonitrile from explicit number of molecules and density, cubic box with auto-determined size")
out, details = packmol(
    molecules=[water, acetonitrile],
    n_molecules=[32, 16],
    density=density,
    return_details=True,
)
printsummary(out, details)
out.write("water-acetonitrile-3.xyz")
view(out, width=300, height=300, padding=-2, picture_path="picture11.png")


print("2-1 water-acetonitrile from explicit number of molecules and box")
out = packmol(
    molecules=[water, acetonitrile],
    n_molecules=[32, 16],
    box_bounds=[0, 0, 0, 13.2, 13.2, 13.2],
)
printsummary(out)
out.write("water-acetonitrile-4.xyz")
view(out, width=300, height=300, padding=-2, picture_path="picture12.png")


print("Experimental feature (AMS2025): guess density for mixture")
print("Note: This density is meant to be equilibrated with NPT MD. It can be very inaccurate!")
out = packmol([water, acetonitrile], mole_fractions=[x_water, x_acetonitrile], n_atoms=100)
print(f"Guessed density: {out.get_density():.2f} kg/m^3")
view(out, width=200, height=200, padding=-2, picture_path="picture13.png")


# ## NaCl solution (solvent with 1 or more solutes)

sodium = from_smiles("[Na+]")
chloride = from_smiles("[Cl-]")


# For dilute solutions, it can be useful to specify the number of solute species, and fill up the rest of the box with solvent.
#
# The required number of solvent molecules are then added to fill up the box to the target overall density or number of atoms.
#
# This feature can be used if exactly **one** of the elements of the ``n_molecules`` list is None.

print(
    "NaCl solution from approximate density (in g/cm^3) and box bounds, and auto-determined number of solvent molecules"
)
out = packmol(
    [sodium, chloride, water],
    n_molecules=[5, 5, None],
    density=1.029,
    box_bounds=[0, 0, 0, 19, 19, 19],
)
printsummary(out)
out.write("sodium-chloride-solution-1.xyz")
view(out, width=300, height=300, padding=-3, fixed_atom_size=False, picture_path="picture14.png")


# Specify the total number of atoms instead of box bounds, and auto-determine a cubic box:

print("NaCl solution from approximate number of atoms and density:")
out = packmol([sodium, chloride, water], n_molecules=[5, 5, None], density=1.029, n_atoms=500)
printsummary(out)
out.write("sodium-chloride-solution-2.xyz")
view(out, width=300, height=300, padding=-3, fixed_atom_size=False, picture_path="picture15.png")


# Specify the total number of atoms instead of the density (less useful option):

print("NaCl solution from approximate number of atoms and box_bounds:")
out = packmol(
    [sodium, chloride, water],
    n_molecules=[5, 5, None],
    n_atoms=500,
    box_bounds=[0, 0, 0, 12, 18, 24],
)
printsummary(out)
out.write("sodium-chloride-solution-3.xyz")
view(out, width=300, height=300, padding=-3, fixed_atom_size=False, picture_path="picture16.png")


# ## Pack inside sphere
#
# Set ``sphere=True`` to pack in a sphere (non-periodic) instead of in a periodic box. The sphere will be centered near the origin.

print("water in a sphere from exact density and number of molecules")
out, details = packmol(molecules=[water], n_molecules=[100], density=1.0, return_details=True, sphere=True)
printsummary(out, details)
print(f"Radius  of sphere: {details['radius']:.3f} ang.")
print(f"Center of mass xyz (ang): {out.get_center_of_mass()}")
out.write("water-sphere.xyz")
view(out, width=300, height=300, padding=-3, picture_path="picture17.png")


print(
    "2-1 water-acetonitrile in a sphere from exact density (in g/cm^3) and "
    "approximate number of atoms and mole fractions"
)
out, details = packmol(
    molecules=[water, acetonitrile],
    mole_fractions=[x_water, x_acetonitrile],
    n_atoms=500,
    density=density,
    return_details=True,
    sphere=True,
)
printsummary(out, details)
out.write("water-acetonitrile-sphere.xyz")
view(out, width=300, height=300, padding=-3, picture_path="picture18.png")


# ## Packing ions, total system charge
#
# The total system charge will be sum of the charges of the constituent molecules.
#
# In PLAMS, ``molecule.properties.charge`` specifies the charge:

ammonium = from_smiles("[NH4+]")  # ammonia.properties.charge == +1
chloride = from_smiles("[Cl-]")  # chloride.properties.charge == -1
print("3 water molecules, 3 ammonium, 1 chloride (non-periodic)")
print("Initial charges:")
print(f"Water: {water.properties.get('charge', 0)}")
print(f"Ammonium: {ammonium.properties.get('charge', 0)}")
print(f"Chloride: {chloride.properties.get('charge', 0)}")
out = packmol(molecules=[water, ammonium, chloride], n_molecules=[3, 3, 1], density=0.4, sphere=True)
tot_charge = out.properties.get("charge", 0)
print(f"Total charge of packmol-generated system: {tot_charge}")
out.write("water-ammonium-chloride.xyz")
view(out, width=300, height=300, picture_path="picture19.png")


# ## Microsolvation
# ``packmol_microsolvation`` can create a microsolvation sphere around a solute.

from scm.plams import packmol_microsolvation

out = packmol_microsolvation(solute=acetonitrile, solvent=water, density=1.5, threshold=4.0)
# for microsolvation it's a good idea to have a higher density than normal to get enough solvent molecules
print(f"Microsolvated structure: {len(out)} atoms.")
out.write("acetonitrile-microsolvated.xyz")

view(out, width=300, height=300, padding=-1, picture_path="picture20.png")


# ## Solid-liquid or solid-gas interfaces
# First, create a slab using the ASE ``fcc111`` function

from scm.plams import plot_molecule, fromASE
from ase.build import fcc111

rotation = "90x,0y,0z"  # sideview of slab
slab = fromASE(fcc111("Al", size=(4, 6, 3), vacuum=15.0, orthogonal=True, periodic=True))
view(slab, width=300, height=300, direction="along_y", fixed_atom_size=False, picture_path="picture21.png")


print("water surrounding an Al slab, from an approximate density")
out = packmol_around(slab, water, density=1.0)
printsummary(out)
out.write("al-water-pure.xyz")
view(out, width=300, height=300, direction="along_y", fixed_atom_size=False, picture_path="picture22.png")


print("2-1 water-acetonitrile mixture surrounding an Al slab, from mole fractions and an approximate density")
out = packmol_around(slab, [water, acetonitrile], mole_fractions=[x_water, x_acetonitrile], density=density)
printsummary(out)
out.write("al-water-acetonitrile.xyz")
view(out, width=300, height=300, direction="along_y", fixed_atom_size=False, picture_path="picture23.png")


from ase.build import surface

print("water surrounding non-orthorhombic Au(211) slab, from an approximate number of molecules")
print("NOTE: non-orthorhombic cell, results are approximate, requires AMS2025")
slab = surface("Au", (2, 1, 1), 6)
slab.center(vacuum=11.0, axis=2)
slab.set_pbc(True)
out = packmol_around(fromASE(slab), [water], n_molecules=[32], tolerance=1.8)
out.write("Au211-water.xyz")
print(f"{out.lattice=}")
view(out, width=300, height=300, direction="along_y", fixed_atom_size=False, picture_path="picture24.png")


# ## Pack inside voids in crystals
#
# Use the ``packmol_around`` function. You can decrease ``tolerance`` if you need to pack very tightly. The default value for ``tolerance`` is 2.0.

from scm.plams import fromASE
from ase.build import bulk

bulk_Al = fromASE(bulk("Al", cubic=True).repeat((3, 3, 3)))
rotation = "-85x,5y,0z"
view(bulk_Al, width=300, height=300, direction="corner_z", fixed_atom_size=False, picture_path="picture25.png")


out = packmol_around(
    current=bulk_Al,
    molecules=[from_smiles("[H]"), from_smiles("[He]")],
    n_molecules=[50, 20],
    tolerance=1.5,
)
printsummary(out)
out.write("al-bulk-with-h-he.xyz")
view(out, width=300, height=300, direction="corner_z", fixed_atom_size=False, picture_path="picture26.png")


# ## Bonds, atom properties (force field types, regions, ...)
#
# The ``packmol()`` function accepts the arguments ``keep_bonds`` and ``keep_atom_properties``. These options will keep the bonds defined for the constitutent molecules, as well as any atomic properties.
#
# The bonds and atom properties are easiest to see by printing the System block for an AMS job:

from scm.plams import Settings

water = from_smiles("O")
n2 = from_smiles("N#N")

# delete properties coming from from_smiles
for at in water:
    at.properties = Settings()
for at in n2:
    at.properties = Settings()

water[1].properties.region = "oxygen_atom"
water[2].properties.mass = 2.014  # deuterium
water.delete_bond(water[1, 2])  # delete bond between atoms 1 and 2 (O and H)


from scm.plams import AMSJob

out = packmol([water, n2], n_molecules=[2, 1], density=0.5)
print(AMSJob(molecule=out).get_input())


# By default, the ``packmol()`` function assigns regions called ``mol0``, ``mol1``, etc. to the different added molecules. The ``region_names`` option lets you set custom names.

out = packmol(
    [water, n2],
    n_molecules=[2, 1],
    density=0.5,
    region_names=["water", "nitrogen_molecule"],
)
print(AMSJob(molecule=out).get_input())


# Below, we also set ``keep_atom_properties=False``, this will remove the previous regions (in this example "oxygen_atom") and mass.

out = packmol([water, n2], n_molecules=[2, 1], density=0.5, keep_atom_properties=False)
print(AMSJob(molecule=out).get_input())


# ``keep_bonds=False`` will additionally ignore any defined bonds:

out = packmol(
    [water, n2],
    n_molecules=[2, 1],
    density=0.5,
    region_names=["water", "nitrogen_molecule"],
    keep_bonds=False,
    keep_atom_properties=False,
)
print(AMSJob(molecule=out).get_input())