# AUTOGENERATED FILE! PLEASE DON'T EDIT
"""This module is for all things related to atoms, molecules and their simulations"""
import k1lib, torch, math
import k1lib.cli as cli
from typing import Dict, List, Union
import matplotlib.pyplot as plt
from collections import deque
from .atom import Atom, _depthAuto
__all__ = ["System"]
def distV(x:torch.Tensor) -> torch.Tensor:
    """Distance vectors of points.

:param x: location Tensor of shape (n, 3)
:returns: vector Tensor of shape (n, n, 3)"""
    n = x.shape[0]
    return x.view(1, n, 3) - x.view(n, 1, 3)
@k1lib.patch(Atom)
def _s_bondLength(self, idx_i:Dict[str, int], bondLengths:torch.Tensor):
    """Calculates bond length for all bonds in this atom and stores in
``bondLengths``."""
    bonds = list(set(self.bonds)) + self.eClouds
    for atom in bonds:
        nBonds = self.nBonds(atom)
        bL = self.radius[nBonds-1] + atom.radius[nBonds-1]
        bondLengths[idx_i[self.idx], idx_i[atom.idx]] = bL
        bondLengths[idx_i[atom.idx], idx_i[self.idx]] = bL
class System:
    def __init__(self, mapping:Dict[str, Atom]):
        """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"""
        self.mapping = mapping; self.atoms = list(mapping.values())
        self.i_idx = list(mapping.keys()); self.n = len(self.i_idx)
        self.idx_i = {idx:i for i, idx in enumerate(self.i_idx)}
        self._bondLengths = torch.zeros(self.n, self.n)
        for atom in self.atoms: atom._s_bondLength(self.idx_i, self._bondLengths)
        self._graphDistances = self._calcGraphDistances()
        self._bondMask = (self._bondLengths > 0) + 0
        self._x = 50*torch.randn(self.n, 3)
        self._v = torch.zeros(self.n, 3)
        self._a = torch.zeros(self.n, 3)
        self.iMask = 1-torch.eye(self.n) # identity mask
    def __getitem__(self, idx:Union[str, int, Atom]) -> Union[Atom, Atom, int]:
        """Look up Atom that has idx (str/int). If Atom is passed in, look
for the index of it."""
        if isinstance(idx, str): return self.mapping[idx]
        if isinstance(idx, int): return self.mapping[self.i_idx[idx]]
        if isinstance(idx, Atom): return self.idx_i[idx.idx]
        raise RuntimeError(f"Don't understand input of type {type(idx)}")
    def x(self, a:Atom) -> torch.Tensor:
        """Get current location vector of the specified :class:`Atom`."""
        return self._x[self.idx_i[a.idx]]
    def v(self, a:Atom):
        """Get current velocity vector of the specified :class:`Atom`."""
        return self._v[self.idx_i[a.idx]]
@k1lib.patch(Atom)
def _s_setGD(self:Atom, s:System, atom:Atom, gd:torch.Tensor, v:int, todos:deque, autoCount):
    """
: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"""
    if v > 4: return
    x = self._s_gi; y = atom._s_gi # int indexes in system
    if v < gd[x, y]:
        autoCount(); gd[y, x] = gd[x, y] = v
        for a in self.uniqueBonds + self.eClouds:
            todos.append((a, self, atom)) # tuple (atom to be updated, (what happens here))
        for a in atom.uniqueBonds + atom.eClouds:
            todos.append((a, atom, self))
@k1lib.patch(System)
def _calcGraphDistances(self):
    gd = torch.ones(self.n, self.n) * float("inf")
    todos = deque(); self._gDAutoCount = k1lib.AutoIncrement()
    for i, a in enumerate(self.atoms): a._s_gi = i # cached i value
    for a in self.atoms: a._s_setGD(self, a, gd, 0, todos, self._gDAutoCount)
    while len(todos) > 0:
        mainA, _as, _ad = todos.popleft() # "source" and "destination"
        mainA._s_setGD(self, _ad, gd, gd[_as._s_gi, _ad._s_gi] + 1, todos, self._gDAutoCount)
    return gd
@k1lib.patch(System)
def _calcForces(self):
    n = self.n; dV = distV(self._x) # (n, n, 3)
    l = torch.sqrt((dV**2).sum(dim=2)) + 0.001
    dV = dV / l.view(n, n, 1) # normalized direction
    aBond = 10 * dV * torch.tanh((l-self._bondLengths)*self._bondMask/100).view(n, n, 1)
    vBond = aBond.sum(dim=1)

    aCou = dV * ((100/l)**4).view(n, n, 1) # coulomb forces on every atom
    vCou = .03 * aCou.clearNan().sum(dim=1)

    aCouC = dV * ((100/l)**2).view(n, n, 1) # close coulomb forces
    couCMask = ((self._graphDistances < 4) + 0).view(n, n, 1)
    vCouC = .3 * (aCouC * couCMask).clearNan().sum(dim=1)

    self._a = 10 * (vBond - vCou - vCouC)
    return aBond, aCou, l, dV, vBond, vCou # return stuff for debugging purposes
@k1lib.patch(System)
def _update(self, dt:float=1.0):
    self._v += dt * self._a; self._a.zero_()
    self._v = torch.clamp(self._v * 0.98, -10, 10) # friction
    self._x += dt * self._v; self._x *= 0.999 # universe wants to quish things down
@k1lib.patch(System)
def _changeDevice(self, device:str="cpu"):
    self._bondLengths = self._bondLengths.to(device)
    self._graphDistances = self._graphDistances.to(device)
    self._bondMask = self._bondMask.to(device)
    self._x = self._x.to(device)
    self._v = self._v.to(device)
    self._a = self._a.to(device)
    self.iMask = self.iMask.to(device)
@k1lib.patch(System)
def simulate(self, t:int=300, dt:float=1.0, recordXs:bool=True, cuda:bool=False) -> List[torch.Tensor]:
    """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."""
    self._changeDevice("cuda" if cuda else "cpu")
    every = k1lib.Every((t // 100) or 1); xs = []
    for i in range(t):
        self._calcForces(); self._update(dt)
        if recordXs and every(): xs.append(self._x.clone())
    return xs
@k1lib.patch(Atom)
def _system(self, mapping:Dict[str, Atom], gDepth:int):
    if self.gDepth >= gDepth: return
    self.gDepth = gDepth; mapping[self.idx] = self
    for atom in self.bonds: atom._system(mapping, gDepth)
    for eCloud in self.eClouds: eCloud._system(mapping, gDepth)
@k1lib.patch(Atom)
def system(self) -> System:
    """Creates a :class:`System` of the molecule this :class:`Atom` is attached
to."""
    mapping = dict(); self._system(mapping, _depthAuto()); return System(mapping)
@k1lib.patch(Atom)
def _s_plot(self, ax, x:torch.Tensor, idx_i:Dict[str, int], H:bool):
    if not H and self.name == "H": return
    nH = len(self.HBonds); HName = "" if nH == 0 else ("H" if nH == 1 else f"H{nH}")
    s = x[idx_i[self.idx]]
    ax.text(s[0], s[1], s[2], f"{self.name + ('' if H else HName)}", ha="center", va="center")
    for atom in (self.uniqueBonds if H else self.uniqueNonHBonds):
        a = x[idx_i[atom.idx]]
        ax.plot([s[0], a[0]], [s[1], a[1]], [s[2], a[2]])
@k1lib.patch(System)
def plot(self, x:torch.Tensor=None, ax:"matplotlib.axes.Axes"=None, H:bool=True):
    """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."""
    if x is None: x = self._x
    if ax is None: ax = plt.figure(dpi=150).add_subplot(projection="3d")
    ax.clear(); com = x.mean(dim=0) # center of mass
    box = math.sqrt(((x - com.view(1, -1))**2).sum(dim=1).max())
    ax.set_xlim(com[0]-box, com[0]+box)
    ax.set_ylim(com[1]-box, com[1]+box)
    ax.set_zlim(com[2]-box, com[2]+box)
    for atom in self.atoms:
        if atom.name != "_e": atom._s_plot(ax, x, self.idx_i, H)
    return ax
@k1lib.patch(System)
def animate(self, xs:List[torch.Tensor], rotateSpeed=0.5, H:bool=True):
    """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."""
    fig = plt.figure(); ax = fig.add_subplot(projection="3d"); plt.close(fig)
    def update(frame):
        self.plot(xs[frame], ax, H); ax.view_init(frame*rotateSpeed, frame*2*rotateSpeed)
    a = k1lib.viz.FAnim(fig, update, len(xs)); plt.close(fig); return a