# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Energy-Derivative Training Contract (Forces, Stress, Charge Gradients)
======================================================================

This example demonstrates the recommended way to obtain forces, virial/stress,
and charge gradients from PME (and Ewald) for machine-learning interatomic
potential (MLIP) training: **energy is the only differentiable output, and all
derivatives are taken from it with** ``torch.autograd.grad``.

The legacy ``compute_forces`` / ``compute_virial`` / ``compute_charge_gradients``
/ ``hybrid_forces`` flags on ``particle_mesh_ewald`` and ``ewald_summation`` are
deprecated and emit a ``DeprecationWarning``. They remain available for
compatibility in v0.4.0, so this example also checks that the autograd results
match the legacy direct outputs, which is the migration check.

Monopole Ewald/PME/slab entry points accept ``energy_reduction="atom"`` (default)
or ``"system"``. Atom mode returns per-atom energies ``(N,)``; system mode
returns per-system totals ``(B,)`` for batched per-system losses. Direct-output
fields (forces, charge gradients, virials) keep their existing shapes. On Torch
CUDA, eager atom mode may synchronize when proving a materialized uniform
cotangent; ``energy_reduction="system"`` is structurally sync-free.

In this example you will learn:

- Forces from energy autograd: ``F = -grad(E.sum(), positions)``
- Force-loss training (double-backward) with ``create_graph=True``
- Geometry-dependent charges ``q(R)``: the full force includes the
  ``dE/dq . dq/dR`` charge-model chain-rule term
- Strain-first virial/stress and stress-loss training
- Charge gradients from energy autograd: ``dE/dq = grad(E.sum(), charges)``
- ``energy_reduction="system"`` for batched per-system energy layout ``(B,)``

.. important::
    This script is intended as an API demonstration. Do not use this script
    for performance benchmarking; refer to the `benchmarks` folder instead.
"""

# %%
# Setup and Imports
# -----------------
# We use ``particle_mesh_ewald`` throughout; ``ewald_summation`` follows the
# exact same energy-derivative contract.

from __future__ import annotations

import warnings

import numpy as np
import torch

from nvalchemiops.torch.interactions.electrostatics import particle_mesh_ewald
from nvalchemiops.torch.neighbors import neighbor_list as neighbor_list_fn

# %%
# Configure Device
# ----------------

if torch.cuda.is_available():
    device = torch.device("cuda:0")
    print("Using CUDA device")
    print(f"  {torch.cuda.get_device_name(0)}")
else:
    device = torch.device("cpu")
    print("Using CPU")

dtype = torch.float64

# %%
# Create a NaCl Crystal System
# ----------------------------
# A small NaCl rock-salt supercell (2x2x2 = 16 atoms) is enough to demonstrate
# every derivative path. Charges use the same floating dtype as the geometry.


def create_nacl_system(n_cells: int = 2, lattice_constant: float = 5.64):
    """Create a NaCl crystal supercell."""
    base_positions = np.array([[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]])
    base_charges = np.array([1.0, -1.0])

    positions = []
    charges = []
    for i in range(n_cells):
        for j in range(n_cells):
            for k in range(n_cells):
                offset = np.array([i, j, k])
                for pos, charge in zip(base_positions, base_charges):
                    positions.append((pos + offset) * lattice_constant)
                    charges.append(charge)

    positions = torch.tensor(np.array(positions), dtype=dtype, device=device)
    charges = torch.tensor(np.array(charges), dtype=dtype, device=device)
    cell = torch.eye(3, dtype=dtype, device=device) * lattice_constant * n_cells
    cell = cell.unsqueeze(0)
    pbc = torch.tensor([[True, True, True]], dtype=torch.bool, device=device)
    return positions, charges, cell, pbc


positions, charges, cell, pbc = create_nacl_system(n_cells=2)

# Rattle the perfect lattice so the per-atom forces are nonzero -- otherwise the
# centrosymmetric crystal gives zero forces and the autograd-vs-direct force
# check below would compare 0 to 0.
torch.manual_seed(0)
positions = positions + 0.05 * torch.randn_like(positions)

print(f"\nSystem: {len(positions)} atoms NaCl crystal (rattled)")

# Build the real-space neighbor list once; reuse it for every call below.
neighbor_list, neighbor_ptr, neighbor_shifts = neighbor_list_fn(
    positions, 8.0, cell=cell, pbc=pbc, return_neighbor_list=True
)

# Fixed PME parameters keep the example deterministic and fast.
pme_kwargs = dict(
    alpha=0.35,
    mesh_dimensions=(32, 32, 32),
    spline_order=4,
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
)

# %%
# Forces From Energy Autograd
# ---------------------------
# With no deprecated flag, the call returns the per-atom energy only. The full
# force is the negative gradient of the total energy w.r.t. positions.

positions_f = positions.detach().requires_grad_(True)
energy = particle_mesh_ewald(positions_f, charges, cell, **pme_kwargs)
print(f"\nenergy shape: {tuple(energy.shape)}  (per-atom, energy_reduction='atom')")

forces = -torch.autograd.grad(energy.sum(), positions_f)[0]
print(f"forces shape: {tuple(forces.shape)}")
print(f"max force magnitude: {forces.norm(dim=1).max().item():.6f}")

# %%
# Migration check: the autograd force equals the legacy ``compute_forces=True``
# direct output (the deprecated flag still works, with a ``DeprecationWarning``).

with warnings.catch_warnings():
    warnings.filterwarnings(
        "ignore",
        category=DeprecationWarning,
        message=".*compute_forces.*",
    )
    _, forces_direct = particle_mesh_ewald(
        positions, charges, cell, compute_forces=True, **pme_kwargs
    )

force_diff = (forces - forces_direct).abs().max().item()
print(f"max |autograd force - direct force|: {force_diff:.2e}")
if force_diff >= 1e-6:
    raise RuntimeError(f"autograd force does not match direct force: {force_diff:.2e}")

# %%
# Force-Loss Training (Double-Backward)
# -------------------------------------
# To train on a force loss, build the force with ``create_graph=True`` so the
# subsequent ``loss.backward()`` can differentiate through the force. We use a
# trivial linear charge model ``q(R)`` to stand in for a learned charge head.

torch.manual_seed(0)
weight = torch.randn(3, dtype=dtype, device=device, requires_grad=True)


def charge_model(pos: torch.Tensor) -> torch.Tensor:
    """Toy geometry-dependent charge model with enforced neutrality."""
    raw = pos @ weight
    return raw - raw.mean()


positions_q = positions.detach().requires_grad_(True)
charges_qr = charge_model(positions_q)  # kept in the graph -> full q(R) force

energy = particle_mesh_ewald(positions_q, charges_qr, cell, **pme_kwargs)
forces_qr = -torch.autograd.grad(energy.sum(), positions_q, create_graph=True)[0]

target_forces = torch.zeros_like(forces_qr)
force_loss = (forces_qr - target_forces).pow(2).sum()
force_loss.backward()  # differentiates through the force construction

print(f"\nq(R) force-loss backward OK; weight.grad shape: {tuple(weight.grad.shape)}")
if not torch.isfinite(weight.grad).all():
    raise RuntimeError("q(R) force-loss produced non-finite parameter gradients")

# %%
# Because ``charges = charge_model(positions)`` stays connected to ``positions``,
# the autograd force includes both the fixed-charge term and the
# ``dE/dq . dq/dR`` chain-rule term. The legacy ``compute_forces=True`` output is
# only the fixed-charge partial and does not include the charge-model term --
# this is the central reason direct force output on the full API is deprecated.

# %%
# Strain-First Virial and Stress
# ------------------------------
# ``strain`` is not a PME argument. Build a differentiable strain tensor, deform
# positions and cell by ``I + strain``, and let autograd map gradients back to
# strain. The virial is ``W = -dE/d(strain)`` and tensile-positive stress is
# ``dE/d(strain) / V``; see ``conventions.md`` for the project-wide sign.

num_systems = cell.shape[0]
positions_s = positions.detach().requires_grad_(True)
strain = torch.zeros(num_systems, 3, 3, device=device, dtype=dtype, requires_grad=True)
eye = torch.eye(3, device=device, dtype=dtype).unsqueeze(0)
deform = eye + strain

# Single system: every atom maps to system 0.
batch_idx = torch.zeros(positions_s.shape[0], dtype=torch.int32, device=device)
positions_def = torch.einsum("ni,nij->nj", positions_s, deform[batch_idx])
cell_def = torch.einsum("bij,bjk->bik", cell, deform)

energy = particle_mesh_ewald(positions_def, charges, cell_def, **pme_kwargs)
grad_strain = torch.autograd.grad(energy.sum(), strain)[0]
virial = -grad_strain
volume = torch.abs(torch.linalg.det(cell_def))
stress = grad_strain / volume[:, None, None]  # tensile-positive Cauchy

print(f"\nvirial shape: {tuple(virial.shape)}")
print(f"stress shape: {tuple(stress.shape)}")

# %%
# Migration check: the strain-first virial equals the legacy
# ``compute_virial=True`` direct virial (both are ``-dE/d(strain)``).

with warnings.catch_warnings():
    warnings.filterwarnings(
        "ignore",
        category=DeprecationWarning,
        message=".*compute_virial.*",
    )
    _, virial_direct = particle_mesh_ewald(
        positions, charges, cell, compute_virial=True, **pme_kwargs
    )

virial_diff = (virial - virial_direct).abs().max().item()
print(f"max |strain-first virial - direct virial|: {virial_diff:.2e}")
if virial_diff >= 1e-5:
    raise RuntimeError(
        f"strain-first virial does not match direct virial: {virial_diff:.2e}"
    )

# %%
# Stress-Loss Training (Double-Backward)
# --------------------------------------
# Stress training uses the same strain-first recipe with ``create_graph=True``,
# so the stress loss back-propagates to model parameters.

weight_s = torch.randn(3, dtype=dtype, device=device, requires_grad=True)
positions_s = positions.detach().requires_grad_(True)
strain = torch.zeros(num_systems, 3, 3, device=device, dtype=dtype, requires_grad=True)
deform = torch.eye(3, device=device, dtype=dtype).unsqueeze(0) + strain
positions_def = torch.einsum("ni,nij->nj", positions_s, deform[batch_idx])
cell_def = torch.einsum("bij,bjk->bik", cell, deform)

charges_qr = positions_def @ weight_s
charges_qr = charges_qr - charges_qr.mean()

energy = particle_mesh_ewald(positions_def, charges_qr, cell_def, **pme_kwargs)
grad_strain = torch.autograd.grad(energy.sum(), strain, create_graph=True)[0]
virial = -grad_strain
volume = torch.abs(torch.linalg.det(cell_def))
stress = grad_strain / volume[:, None, None]

stress_loss = stress.pow(2).sum()
stress_loss.backward()

stress_grad_ok = bool(torch.isfinite(weight_s.grad).all())
print(f"\nstress-loss backward OK; weight_s.grad finite: {stress_grad_ok}")
if not stress_grad_ok:
    raise RuntimeError("stress-loss produced non-finite parameter gradients")

# %%
# Combined Energy + Force + Stress Loss (Performance Tip)
# -------------------------------------------------------
# When a single loss mixes energy, forces, AND stress, take the forces and the
# virial from **one** ``torch.autograd.grad`` call over both ``positions`` and
# ``strain`` -- NOT two separate calls. Each ``create_graph=True`` ``grad`` call
# builds its own first-derivative graph node, and ``loss.backward()`` then runs the
# (O(K*N)) reciprocal second-derivative once per node. Fusing the two into a single
# call avoids duplicate reciprocal double-backward work; the gradients are
# bit-identical to the two-call form.
weight_m = torch.randn(3, dtype=dtype, device=device, requires_grad=True)
positions_m = positions.detach().requires_grad_(True)
strain_m = torch.zeros(
    num_systems, 3, 3, device=device, dtype=dtype, requires_grad=True
)
deform_m = torch.eye(3, device=device, dtype=dtype).unsqueeze(0) + strain_m
positions_md = torch.einsum("ni,nij->nj", positions_m, deform_m[batch_idx])
cell_md = torch.einsum("bij,bjk->bik", cell, deform_m)
charges_m = positions_md @ weight_m
charges_m = charges_m - charges_m.mean()

energy = particle_mesh_ewald(positions_md, charges_m, cell_md, **pme_kwargs)
# ONE combined grad call -> one double-backward (do this instead of separate
# ``grad(E, positions)`` and ``grad(E, strain)`` calls):
grad_pos, grad_strain = torch.autograd.grad(
    energy.sum(), (positions_m, strain_m), create_graph=True
)
# ``positions_m`` are the undeformed reference coordinates, so ``forces_m`` are
# reference-frame forces. Differentiate with respect to ``positions_md`` instead
# when training against deformed-coordinate force targets.
forces_m = -grad_pos
virial_m = -grad_strain
volume_m = torch.abs(torch.linalg.det(cell_md))
stress_m = grad_strain / volume_m[:, None, None]

mixed_loss = energy.sum() + forces_m.pow(2).sum() + stress_m.pow(2).sum()
mixed_loss.backward()
mixed_grad_ok = bool(torch.isfinite(weight_m.grad).all())
print(f"\nmixed E+F+stress backward OK; weight_m.grad finite: {mixed_grad_ok}")
if not mixed_grad_ok:
    raise RuntimeError("mixed-loss produced non-finite parameter gradients")

# %%
# Charge Gradients From Energy Autograd
# -------------------------------------
# ``dE/dq`` is an ordinary gradient of the energy w.r.t. charges.

charges_g = charges.detach().requires_grad_(True)
energy = particle_mesh_ewald(positions, charges_g, cell, **pme_kwargs)
charge_grad = torch.autograd.grad(energy.sum(), charges_g)[0]
print(f"\ncharge gradient shape: {tuple(charge_grad.shape)}")

with warnings.catch_warnings():
    warnings.filterwarnings(
        "ignore",
        category=DeprecationWarning,
        message=".*compute_forces.*",
    )
    _, _, charge_grad_direct = particle_mesh_ewald(
        positions,
        charges,
        cell,
        compute_forces=True,
        compute_charge_gradients=True,
        **pme_kwargs,
    )

cg_diff = (charge_grad - charge_grad_direct).abs().max().item()
print(f"max |autograd dE/dq - direct dE/dq|: {cg_diff:.2e}")
if cg_diff >= 1e-6:
    raise RuntimeError(f"autograd dE/dq does not match direct dE/dq: {cg_diff:.2e}")

# %%
# Summary
# -------
# This example demonstrated the energy-derivative training contract:
#
# 1. **Forces** -- ``F = -torch.autograd.grad(E.sum(), positions)[0]``; with
#    ``create_graph=True`` for force-loss training.
# 2. **q(R) forces** -- keep ``charges = charge_model(positions)`` in the graph so
#    the full ``dE/dR`` includes the ``dE/dq . dq/dR`` charge-model term.
# 3. **Virial / stress** -- strain-first: deform positions and cell by
#    ``I + strain``, then ``virial = -grad(E.sum(), strain)`` and
#    ``stress = grad(E.sum(), strain) / volume[:, None, None]``.
# 4. **Charge gradients** -- ``dE/dq = torch.autograd.grad(E.sum(), charges)[0]``.
# 5. **Combined E + F + stress loss** -- take forces and virial from a single
#    ``torch.autograd.grad(E.sum(), (positions, strain), create_graph=True)`` call,
#    not two separate calls, so the reciprocal double-backward runs once.
# 6. **Batched per-system layout** -- pass ``energy_reduction="system"`` to get
#    ``(B,)`` energies for per-system losses without manual ``scatter_add``.
#
# Each autograd result matched the corresponding (deprecated) direct kernel
# output, confirming the migration is numerically exact.

print("\nEnergy-derivative training contract example complete!")