# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
"""This module is for all things related to atoms, molecules and their simulations"""
import k1lib, math, k1lib.cli as cli
from typing import Dict, List, Union
plt = k1lib.dep("matplotlib.pyplot")
from collections import deque
from .atom import Atom, _depthAuto
try: import torch; hasTorch = True
except: torch = k1lib.Object().withAutoDeclare(lambda: type("RandomClass", (object, ), {})); hasTorch = False
__all__ = ["System"]
def distV(x:torch.Tensor) -> torch.Tensor:                                       # distV
    """Distance vectors of points.

:param x: location Tensor of shape (n, 3)
:returns: vector Tensor of shape (n, n, 3)"""                                    # distV
    n = x.shape[0]                                                               # distV
    return x.view(1, n, 3) - x.view(n, 1, 3)                                     # distV
@k1lib.patch(Atom)                                                               # distV
def _s_bondLength(self, idx_i:Dict[str, int], bondLengths:torch.Tensor):         # _s_bondLength
    """Calculates bond length for all bonds in this atom and stores in
``bondLengths``."""                                                              # _s_bondLength
    bonds = list(set(self.bonds)) + self.eClouds                                 # _s_bondLength
    for atom in bonds:                                                           # _s_bondLength
        nBonds = self.nBonds(atom)                                               # _s_bondLength
        bL = self.radius[nBonds-1] + atom.radius[nBonds-1]                       # _s_bondLength
        bondLengths[idx_i[self.idx], idx_i[atom.idx]] = bL                       # _s_bondLength
        bondLengths[idx_i[atom.idx], idx_i[self.idx]] = bL                       # _s_bondLength
class System:                                                                    # System
    def __init__(self, mapping:Dict[str, Atom]):                                 # System
        """Creates a new system that contains a molecule so that it can be
simulated. Not intended to be used by the end user. Use :meth:`Atom.system`
instead.

:param mapping: maps from atom index to :class:`Atom` object"""                  # System
        self.mapping = mapping; self.atoms = list(mapping.values())              # System
        self.i_idx = list(mapping.keys()); self.n = len(self.i_idx)              # System
        self.idx_i = {idx:i for i, idx in enumerate(self.i_idx)}                 # System
        self._bondLengths = torch.zeros(self.n, self.n)                          # System
        for atom in self.atoms: atom._s_bondLength(self.idx_i, self._bondLengths) # System
        self._graphDistances = self._calcGraphDistances()                        # System
        self._bondMask = (self._bondLengths > 0) + 0                             # System
        self._x = 50*torch.randn(self.n, 3)                                      # System
        self._v = torch.zeros(self.n, 3)                                         # System
        self._a = torch.zeros(self.n, 3)                                         # System
        self.iMask = 1-torch.eye(self.n) # identity mask                         # System
    def __getitem__(self, idx:Union[str, int, Atom]) -> Union[Atom, Atom, int]:  # System
        """Look up Atom that has idx (str/int). If Atom is passed in, look
for the index of it."""                                                          # System
        if isinstance(idx, str): return self.mapping[idx]                        # System
        if isinstance(idx, int): return self.mapping[self.i_idx[idx]]            # System
        if isinstance(idx, Atom): return self.idx_i[idx.idx]                     # System
        raise RuntimeError(f"Don't understand input of type {type(idx)}")        # System
    def x(self, a:Atom) -> torch.Tensor:                                         # System
        """Get current location vector of the specified :class:`Atom`."""        # System
        return self._x[self.idx_i[a.idx]]                                        # System
    def v(self, a:Atom):                                                         # System
        """Get current velocity vector of the specified :class:`Atom`."""        # System
        return self._v[self.idx_i[a.idx]]                                        # System
@k1lib.patch(Atom)                                                               # System
def _s_setGD(self:Atom, s:System, atom:Atom, gd:torch.Tensor, v:int, todos:deque, autoCount): # _s_setGD
    """
:param a: other (constant) atom that I'm trying to get distance to
:param gd: graph distances
:param v: value (aka distance)
:param todos: queue of atoms to be processed
:param autoCount: to measure perf"""                                             # _s_setGD
    if v > 4: return                                                             # _s_setGD
    x = self._s_gi; y = atom._s_gi # int indexes in system                       # _s_setGD
    if v < gd[x, y]:                                                             # _s_setGD
        autoCount(); gd[y, x] = gd[x, y] = v                                     # _s_setGD
        for a in self.uniqueBonds + self.eClouds:                                # _s_setGD
            todos.append((a, self, atom)) # tuple (atom to be updated, (what happens here)) # _s_setGD
        for a in atom.uniqueBonds + atom.eClouds:                                # _s_setGD
            todos.append((a, atom, self))                                        # _s_setGD
@k1lib.patch(System)                                                             # _s_setGD
def _calcGraphDistances(self):                                                   # _calcGraphDistances
    gd = torch.ones(self.n, self.n) * float("inf")                               # _calcGraphDistances
    todos = deque(); self._gDAutoCount = k1lib.AutoIncrement()                   # _calcGraphDistances
    for i, a in enumerate(self.atoms): a._s_gi = i # cached i value              # _calcGraphDistances
    for a in self.atoms: a._s_setGD(self, a, gd, 0, todos, self._gDAutoCount)    # _calcGraphDistances
    while len(todos) > 0:                                                        # _calcGraphDistances
        mainA, _as, _ad = todos.popleft() # "source" and "destination"           # _calcGraphDistances
        mainA._s_setGD(self, _ad, gd, gd[_as._s_gi, _ad._s_gi] + 1, todos, self._gDAutoCount) # _calcGraphDistances
    return gd                                                                    # _calcGraphDistances
@k1lib.patch(System)                                                             # _calcGraphDistances
def _calcForces(self):                                                           # _calcForces
    n = self.n; dV = distV(self._x) # (n, n, 3)                                  # _calcForces
    l = torch.sqrt((dV**2).sum(dim=2)) + 0.001                                   # _calcForces
    dV = dV / l.view(n, n, 1) # normalized direction                             # _calcForces
    aBond = 10 * dV * torch.tanh((l-self._bondLengths)*self._bondMask/100).view(n, n, 1) # _calcForces
    vBond = aBond.sum(dim=1)                                                     # _calcForces
                                                                                 # _calcForces
    aCou = dV * ((100/l)**4).view(n, n, 1) # coulomb forces on every atom        # _calcForces
    vCou = .03 * aCou.clearNan().sum(dim=1)                                      # _calcForces
                                                                                 # _calcForces
    aCouC = dV * ((100/l)**2).view(n, n, 1) # close coulomb forces               # _calcForces
    couCMask = ((self._graphDistances < 4) + 0).view(n, n, 1)                    # _calcForces
    vCouC = .3 * (aCouC * couCMask).clearNan().sum(dim=1)                        # _calcForces
                                                                                 # _calcForces
    self._a = 10 * (vBond - vCou - vCouC)                                        # _calcForces
    return aBond, aCou, l, dV, vBond, vCou # return stuff for debugging purposes # _calcForces
@k1lib.patch(System)                                                             # _calcForces
def _update(self, dt:float=1.0):                                                 # _update
    self._v += dt * self._a; self._a.zero_()                                     # _update
    self._v = torch.clamp(self._v * 0.98, -10, 10) # friction                    # _update
    self._x += dt * self._v; self._x *= 0.999 # universe wants to quish things down # _update
@k1lib.patch(System)                                                             # _update
def _changeDevice(self, device:str="cpu"):                                       # _changeDevice
    self._bondLengths = self._bondLengths.to(device)                             # _changeDevice
    self._graphDistances = self._graphDistances.to(device)                       # _changeDevice
    self._bondMask = self._bondMask.to(device)                                   # _changeDevice
    self._x = self._x.to(device)                                                 # _changeDevice
    self._v = self._v.to(device)                                                 # _changeDevice
    self._a = self._a.to(device)                                                 # _changeDevice
    self.iMask = self.iMask.to(device)                                           # _changeDevice
@k1lib.patch(System)                                                             # _changeDevice
def simulate(self, t:int=300, dt:float=1.0, recordXs:bool=True, cuda:bool=False) -> List[torch.Tensor]: # simulate
    """Simulate system for ``t`` steps.

:param t: simulation total steps. 100 to view dynamics closely, 1000 to make
    sure it converges, default 300 is sweet spot
:param dt: simulation time between steps
:param recordXs: whether to record locations while the simulation happens. Put
    False to save memory/performance.
:param cuda: if True, do the simulation on the graphics card
:return: if ``recordXs`` is True, returns max 100 location Tensors. The Tensors
    will have shape of (n, 3), where n is the number of atoms and electron clouds
    your molecule has."""                                                        # simulate
    self._changeDevice("cuda" if cuda else "cpu")                                # simulate
    every = k1lib.Every((t // 100) or 1); xs = []                                # simulate
    for i in range(t):                                                           # simulate
        self._calcForces(); self._update(dt)                                     # simulate
        if recordXs and every(): xs.append(self._x.clone())                      # simulate
    return xs                                                                    # simulate
@k1lib.patch(Atom)                                                               # simulate
def _system(self, mapping:Dict[str, Atom], gDepth:int):                          # _system
    if self.gDepth >= gDepth: return                                             # _system
    self.gDepth = gDepth; mapping[self.idx] = self                               # _system
    for atom in self.bonds: atom._system(mapping, gDepth)                        # _system
    for eCloud in self.eClouds: eCloud._system(mapping, gDepth)                  # _system
@k1lib.patch(Atom)                                                               # _system
def system(self) -> System:                                                      # system
    """Creates a :class:`System` of the molecule this :class:`Atom` is attached
to."""                                                                           # system
    mapping = dict(); self._system(mapping, _depthAuto()); return System(mapping) # system
@k1lib.patch(Atom)                                                               # system
def _s_plot(self, ax, x:torch.Tensor, idx_i:Dict[str, int], H:bool):             # _s_plot
    if not H and self.name == "H": return                                        # _s_plot
    nH = len(self.HBonds); HName = "" if nH == 0 else ("H" if nH == 1 else f"H{nH}") # _s_plot
    s = x[idx_i[self.idx]]                                                       # _s_plot
    ax.text(s[0], s[1], s[2], f"{self.name + ('' if H else HName)}", ha="center", va="center") # _s_plot
    for atom in (self.uniqueBonds if H else self.uniqueNonHBonds):               # _s_plot
        a = x[idx_i[atom.idx]]                                                   # _s_plot
        ax.plot([s[0], a[0]], [s[1], a[1]], [s[2], a[2]])                        # _s_plot
@k1lib.patch(System)                                                             # _s_plot
def plot(self, x:torch.Tensor=None, ax:"matplotlib.axes.Axes"=None, H:bool=True): # plot
    """Plots the molecule.
Example::

    %matplotlib widget
    s = mo.CH4(mo.H2O).system()
    s.simulate(); s.plot()

The first line is so that you can rotate the molecule around in 3d interactively.
The 3rd line includes a simulation step, to get the molecule into roughly the
right position.

:param x: location Tensor of shape (n, 3). If none provided, will use current locations
:param ax: Axes object, in case you want to plot this on an existing plot
:param H: whether to include Hydrogens as separate atoms or bunched together to
    another main atom."""                                                        # plot
    if x is None: x = self._x                                                    # plot
    if ax is None: ax = plt.figure(dpi=150).add_subplot(projection="3d")         # plot
    ax.clear(); com = x.mean(dim=0) # center of mass                             # plot
    box = math.sqrt(((x - com.view(1, -1))**2).sum(dim=1).max())                 # plot
    ax.set_xlim(com[0]-box, com[0]+box)                                          # plot
    ax.set_ylim(com[1]-box, com[1]+box)                                          # plot
    ax.set_zlim(com[2]-box, com[2]+box)                                          # plot
    for atom in self.atoms:                                                      # plot
        if atom.name != "_e": atom._s_plot(ax, x, self.idx_i, H)                 # plot
    return ax                                                                    # plot
@k1lib.patch(System)                                                             # plot
def animate(self, xs:List[torch.Tensor], rotateSpeed=0.5, H:bool=True):          # animate
    """Animates the molecule. This requires location information from
:meth:`simulate`. Example::

    s = mo.CH4(mo.H2O).system()
    s.animate(s.simulate())

:param H: whether to include Hydrogens as separate atoms or bunched together to
    another main atom."""                                                        # animate
    fig = plt.figure(); ax = fig.add_subplot(projection="3d"); plt.close(fig)    # animate
    def update(frame):                                                           # animate
        self.plot(xs[frame], ax, H); ax.view_init(frame*rotateSpeed, frame*2*rotateSpeed) # animate
    a = k1lib.viz.FAnim(fig, update, len(xs)); plt.close(fig); return a          # animate
