# NEC Software Library Academic Use Licence Terms (the 'Terms')

# THESE TERMS APPLY TO THE LICENSEE'S USE OF THE "SOFTWARE" AS IDENTIFIED BELOW 
# (INCLUDING ANY PORTION OR ELEMENT OF THE SOFTWARE), INCLUDING ANY UPDATES TO 
# THE SOFTWARE RECEIVED FROM NEC, AND BY DOWNLOADING OR USING THE SOFTWARE (OR 
# RECEIVING ANY UPDATES TO THE SOFTWARE) THE LICENSEE AGREES TO BE BOUND BY THESE 
# TERMS. 

# THESE TERMS ARE A CONTRACTUAL AGREEMENT BETWEEN THE LICENSEE AND NEC. IF THE 
# LICENSEE DOES NOT AGREE TO THESE TERMS, IT CANNOT DOWNLOAD THE SOFTWARE OR USE 
# THE SOFTWARE FOR ANY PURPOSE.


# 1.	Software licence rights for Academic Use

# 1.1	The Licensee is granted a personal, non-exclusive, worldwide, 
# non-transferrable, non-sublicensable, royalty-free and limited licence to use, 
# reproduce, copy, create derivate works of and/or make modifications to the 
# Software solely for Academic Use; and to use and copy the Documentation solely 
# in connection with the rights granted in the foregoing.

# 1.2	The Licensee acknowledges that, unless otherwise expressly agreed in 
# writing by NEC, the licence granted under clause 1.1 does not permit the 
# commercial use of the Software for any purpose, including any distribution or 
# resale of the Software (either directly, or indirectly as part of any package 
# or bundle of software, cloud software or software as a service product), any 
# use in collaborations between academic and commercial entities (including any 
# funding of PhD activities by commercial entities), or any non-commercial use of 
# the Software by a commercial entity.

# 1.3	For the avoidance of doubt, the Licensee is not permitted to use the 
# Software and/or the Documentation under these Terms for any commercial purposes 
# whilst representing that such use is a form of Academic Use. If NEC has the 
# reasonable belief that the Licensee is using the Software and/or the 
# Documentation for any commercial purposes then such use will be deemed a breach 
# of these Terms and NEC may terminate the licence granted under these Terms with 
# immediate effect in accordance with clause 7.2(a).

# 1.4	The Licensee must use the Software and/or the Documentation in 
# accordance with any applicable laws and regulations, including any applicable 
# import and export control laws, associated embargo and sanction rules and 
# regulations relating to the export of software, or materials or products in 
# connection with the export of software. For clarity, the parties acknowledge 
# that the export control laws of United States and the export/import control 
# laws of other countries may govern the use of the Software and/or the 
# Documentation. The Licensee therefore agrees to comply fully with all relevant 
# export laws and regulations, including those of Japan, European Union (and 
# European Union member states), United Kingdom, and the United States, to ensure 
# that the Software and/or the Documentation or any portion thereof is exported, 
# directly or indirectly, in violation of export laws and regulations. This 
# clause will survive the termination or expiration of these Terms.

# 1.5	The Licensee shall ensure that only its Representatives use the 
# Software and that such use is at all times in accordance with these Terms. The 
# Licensee shall procure that each Representative is aware of, and complies with, 
# the obligations and restrictions imposed on the Licensee under these Terms. 

# 1.6	The Licensee shall be liable for the acts and omissions of its 
# Representatives as if they were its own, and shall not provide access to the 
# Software (or permit access by) anyone other than a Representative.

# 1.7	If the Licensee wishes to use the Software, the Documentation and/or 
# any Derivative Works for commercial purposes then it must enter into a separate 
# commercial licence agreement with NEC, on such terms as NEC acting in its sole 
# discretion requires from time to time. For the avoidance of doubt, NEC reserves 
# the right to not enter into any such commercial licence agreement or other 
# arrangements with the Licensee or any third party.


# 2.	Licence rights regarding Derivative Works

# 2.1	If the Licensee develops any derivate works of the Software, any part 
# of the Software and/or any modifications to the Software, either as the result 
# of a pull request received by NEC from the Licensee or a Representative or 
# otherwise (each a 'Derivative Work') or any product or service that uses or 
# incorporates any Derivative Works then:
# (a)	subject to clause 5.1, as between the Licensee and NEC, all 
# Intellectual Property Rights in the Derivative Works belong to and shall remain 
# vested in the Licensee;
# (b)	the Licensee hereby grants to NEC a perpetual, fully paid-up, 
# non-exclusive, royalty-free, transferrable, sublicensable, irrevocable licence 
# to use, copy, modify, reproduce and/or create derivate works of the Derivative 
# Works for any purpose; 
# (c)	the Licensee must provide NEC a copy of the source code of any 
# Derivative Works via its code repository located at https://github.com/nec-research/DIMOS;
# (d)	the Licensee must provide a copy of these Terms with the Derivative 
# Work(s);
# (e)	the Licensee must prominently display "Built using NEC technology" on 
# any related software description, website, user interface, about page, software 
# specification or other technical documentation; and
# (f)	the Licensee must include in any notice text file distributed as part 
# of any Derivate Work: "[Derivate Work] is licensed under the NEC Software 
# Library Academic License. Copyright © 2025 NEC Laboratories Europe GmbH. All 
# Rights Reserved. 

# 2.2	The Licensee is only permitted to use any Derivative Works for Academic 
# Use. The Licensee is not permitted to use any Derivative Works for commercial 
# use, or to permit or otherwise enable any third party to use any Derivative 
# Works for commercial use. Section 1 of the Terms do accordingly apply. 


# 3.	Acceptable use policy

# 3.1	As a condition of use, the Licensee agrees not to use the Software, nor 
# permit it to be used by any Representative or otherwise: 
# (a)	for any purpose that is unlawful under any applicable law or prohibited 
# by these Terms; 
# (b)	to commit any act of fraud or money laundering or otherwise be used in 
# any activities that are deceptive or harmful to others; 
# (c)	to create, store, access, transfer to any third party or otherwise 
# distribute any code or device intended to interfere with, or having the effect 
# of interfering adversely with, the operation of any hardware or software, 
# including any virus, disabling code (including code intended to limit or 
# prevent any use any software or system) or other malicious software (including 
# bugs, worms, logic bombs, malware, trojan horses, ransomware and spyware), or 
# any other material which is unlawful;
# (d)	for the purposes of promoting unsolicited advertising or sending spam, 
# or for the creation or promotion of disinformation;
# (e)	for the purposes of creating, promoting or furthering any defamatory 
# content, including any defamatory statements, images or other media;
# (f)	to simulate communications from any third party service or entity (i.e. 
# phishing communications) in order to collect identity information, 
# authentication credentials, or other information; 
# (g)	in any manner that disrupts the operations, business, equipment, 
# websites or systems of any person or entity (including any denial of service 
# and similar attacks);
# (h)	in any manner that harms or may endanger minors or any other person, or 
# may cause damage or loss to any tangible property or the environment;
# (i)	to gain unauthorised access to or use of any computers, data, systems, 
# accounts or networks of any person;
# (j)	to attempt to circumvent any security controls or mechanisms; 
# (k)	in any manner that engages in or promotes the harassment, abuse, 
# bullying or discrimination of any individuals or groups of individuals; 
# (l)	in any manner that violates, infringes upon, or misappropriates the 
# Intellectual Property Rights or proprietary rights of NEC or any third party; 
# (m)	to sell, publish, rent, lease or otherwise commercialise the Software 
# or any Derivative Work, or assist any third party in relation to the same; or
# (n)	to reverse engineer, decompile or disassemble the Software except and 
# solely to the extent permitted under these Terms.


# 4.	No warranty

# 4.1	The Software and Documentation (including any part thereof and/or any 
# updates provided under these Terms) (together the 'NEC Materials') are provided 
# "as is" and without warranty to the maximum extent permitted by law. To the 
# extent permitted by law, NEC gives no express or implied warranties, 
# guarantees, conditions or obligations under these Terms, including any express 
# or implied terms relating to quality, fitness for any particular purpose, 
# title, non-infringement, or ability to achieve a particular result. 

# 4.2	The Licensee acknowledges that it is solely responsible for determining 
# the appropriateness of using the NEC Materials and assumes any and all risks 
# associated with its use of the NEC Materials, any Derivative Works and/or any 
# other outputs or results relating to any of the foregoing.


# 5.	Intellectual property rights

# 5.1	As between the Licensee and NEC, all rights, including all Intellectual 
# Property Rights, in the NEC Materials belong to and shall remain vested in NEC. 
# Except for the rights expressly granted in these Terms, the Licensee and/or any 
# Representative shall not acquire in any way any title, rights of ownership, or 
# any other Intellectual Property Rights or other rights of any nature in the NEC 
# Materials and no Intellectual Property Rights or other rights are transferred 
# or licensed as a result of these Terms.

# 5.2	The Licensee acknowledges and agrees that certain aspects of the 
# Software provided under these Terms are based on or powered by Open Source 
# Software, either developed by NEC (and which is not based on or derived from 
# Open Source Software licensed to NEC) or derived from, or incorporated into 
# and/or distributed with other Open Source Software licensed to NEC.

# 5.3	NEC may use any feedback and suggestions for improvement relating to 
# the NEC Materials provided by the Licensee or any Representative without charge 
# or limitation ('Feedback'). The Licensee hereby assigns (or shall procure the 
# assignment of) all Intellectual Property Rights in the Feedback with full title 
# guarantee (including by way of present assignment of future Intellectual 
# Property Rights) to NEC at the time such Feedback is first provided to NEC.

# 5.4	The Licensee acknowledges that the Software may collect information 
# about its use of the NEC Materials under these Terms. NEC may use this 
# information to provide and/or improve its services and products. 

# 5.5	If the Licensee (including any Representatives) institutes any claim or 
# other proceedings against NEC or any other entity alleging that any NEC 
# Materials (including any outputs or results), or any part thereof, constitute 
# an infringement of any third party rights, including any Intellectual Property 
# Rights, then any rights granted to the Licensee under these Terms shall 
# immediately terminate.  

# 5.6	The Licensee shall indemnify NEC from and against any losses, claims, 
# damages, liability, costs (including legal and other professional fees) and 
# expenses incurred by NEC as a result of or in connection with any third party 
# action, demand or claim relating to use of the Software and the Documentation 
# if such use is against these Terms or to any transmission, receipt, copying, 
# installation, use, possession or other utilisation of any Derivative Work.


# 6.	Limitation of liability
# 6.1	To the extent permitted by law, NEC shall not be liable to the Licensee 
# (and/or its Representatives) under these Terms for any consequential, indirect 
# or special losses, any loss of profit or revenue; destruction, loss of use or 
# corruption of data; loss or corruption of software or systems; loss or damage 
# to equipment; loss of opportunity; and/or harm to reputation or loss of 
# goodwill.

# 6.2	Subject to clause 6.1 and to the extent permitted by law, NEC's 
# aggregate liability to the Licensee (and/or its Representatives) under these 
# Terms shall not exceed €500.

# 6.3	The rights and remedies provided under these Terms are in addition to, 
# and not exclusive of, any rights or remedies provided by law.


# 7.	Term and termination

# 7.1	These Terms shall commence upon the sooner of the Licensee's: (i) 
# acceptance of these Terms; or (ii) access to or download of any of the NEC 
# Materials, and will continue in full force and effect until terminated in 
# accordance with the provisions of these Terms.

# 7.2	NEC may terminate these Terms with immediate effect without notice if: 
# (a)	the Licensee (including any Representatives) is in breach of any 
# provision of these Terms; 
# (b)	the Licensee takes or has taken against it (other than in relation to a 
# solvent restructuring) any step or action towards its entering bankruptcy, 
# administration, provisional liquidation or any composition or arrangement with 
# its creditors, or it entering into a procedure in any jurisdiction with a 
# similar effect to a procedure listed above; or 
# (c)	the Licensee undergoes a change in beneficial ownership of more than 
# 50% of the issued share capital of the company or a change in the legal power 
# to direct or cause the direction of the general management of the company. 

# 7.3	Upon termination of these Terms for any reason:
# (a)	all rights granted by NEC under these Terms shall immediately cease; 
# (b)	the Licensee shall (and shall ensure that its Representatives shall) 
# cease all use of the NEC Materials;
# (c)	the Licensee shall (and shall procure that each Representatives shall) 
# destroy and delete or, if requested by NEC, return any copies of the Software 
# and/or the Documentation in its possession or control (or in the possession or 
# control of any Representative); and
# (d)	any provision of these Terms that expressly or by implication is 
# intended to continue in force on or after termination or expiry of these Terms 
# shall remain in full force and effect.


# 8.	General

# 8.1	Assignment. The Licensee shall not assign, transfer, mortgage, charge, 
# subcontract, delegate, declare a trust over or deal in any other manner with 
# any of its rights and obligations under these Terms without the prior written 
# consent of NEC. NEC may at any time assign, mortgage, charge, delegate, declare 
# a trust over or deal in any other manner with any or all of its rights under 
# these Terms.

# 8.2	Variation. No variation of these Terms shall be effective unless it is 
# in writing and signed by the parties (or their authorised representatives).

# 8.3	Entire agreement. These Terms constitute the entire agreement between 
# NEC and the Licensee and supersedes and extinguishes all previous agreements 
# and understandings between them, whether written or oral, relating to its 
# subject matter.

# 8.4	No partnership or agency. NEC and the Licensee are independent and are 
# not partners or principal and agent and these Terms does not establish any 
# joint venture, trust, fiduciary or other relationship between them, other than 
# the contractual relationship expressly provided for in it. Neither NEC or the 
# Licensee shall have, nor shall represent that it has, any authority to make any 
# commitments on the other's behalf.

# 8.5	Severance. If any provision of these Terms (or part of any provision) 
# is or becomes illegal, invalid or unenforceable, the legality, validity and 
# enforceability of any other provision of these Terms shall not be affected.

# 8.6	Waiver. A waiver of any right or remedy under these Terms or by any 
# applicable law is only effective if given in writing and shall not be deemed a 
# waiver of any subsequent right or remedy. A failure or delay to exercise any 
# right or remedy provided under these Terms or by law shall not constitute a 
# waiver of that or any other right or remedy, nor shall it prevent or restrict 
# any further exercise of that or any other right or remedy. No single or partial 
# exercise of any right or remedy provided under these Terms or by any applicable 
# law shall prevent or restrict the further exercise of that or any other right 
# or remedy.

# 8.7	Third party rights. Nothing in these Terms confers or purports to 
# confer on any third party any right to enforce any provision of these Terms.

# 8.8	Further assurance. The Licensee (at its own cost) shall, and shall use 
# all reasonable endeavours to procure that any necessary third party shall, 
# promptly execute and deliver such documents and perform such acts as may 
# reasonably be required for the purpose of giving full effect to these Terms.

# 8.9	Governing law and jurisdiction. These Terms and any dispute or claim 
# (including non-contractual disputes or claims) arising out of or in connection 
# with it or its subject matter or formation shall be governed by and construed 
# in accordance with the law of England and Wales. Any dispute or claim arising 
# out of, or in connection with these Terms or its subject matter or formation 
# (including non-contractual disputes or claims) shall be subject to the 
# exclusive jurisdiction of the courts of England and Wales. 


# 9.	Definitions and interpretation
# In these Terms the following words and expressions shall have the following 
# meanings:

# 9.1	"Academic Use" means the use by accredited educational or research 
# institutions for non-commercial purposes, including academic research and/or 
# teaching. 

# 9.2	"Documentation" means description and specification of the Software, or 
# the relevant instructions as to how to use the Software, or other technical 
# manuals or documentation for the Software as made available by NEC from time to 
# time.

# 9.3	"Intellectual Property Rights" means any and all copyright, rights in 
# inventions, patents, know-how, trade secrets, trade marks and trade names, 
# service marks, design rights, rights in get-up, database rights and rights in 
# data, semiconductor chip topography rights, utility models, domain names and 
# all similar rights and, in each case: (a) whether registered or not; (b) 
# including any applications to protect or register such rights; (c) including 
# all renewals and extensions of such rights or applications; (d) whether vested, 
# contingent or future; and (e) wherever existing.

# 9.4	"Licensee" means the legal entity accessing or downloading the Software 
# under these Terms.

# 9.5	"NEC" means NEC Laboratories Europe GmbH, a company incorporated under 
# the laws of Germany with registration number HRB 728558 (VAT No.: DE313703076), 
# being established at Kurfuersten-Anlage 36, 69115 Heidelberg, Germany.

# 9.6	"NEC Software Library" means the software library as available at 
# https://github.com/nec-research/DIMOS.

# 9.7	"Open Source Software" means software which is included or combined 
# with the Software and which is provided to NEC under terms and conditions other 
# than those set out under these Terms, including any ‘open source’ software 
# falling within the Open Source Definition issued by the Open Source Initiative 
# (www.opensource.org/docs/osd) and/or any ‘free software’ as defined by the Free 
# Software Foundation (www.gnu.org/philosophy/free-sw.html).

# 9.8	"Representatives" means, as applicable, the employees, contractors, 
# officers, directors, agents, students of the Licensee and/or any other 
# representatives acting on the Licensee's behalf.

# 9.9	"Software" means the applicable software accessed or downloaded from 
# the NEC Software Library.

# 9.10	Clause and paragraph headings shall not affect the interpretation of 
# these Terms.

# 9.11	A reference to legislation or a legislative provision is a reference to 
# it as amended, extended or re-enacted from time to time, and shall include all 
# subordinate legislation made from time to time under that legislation.

# 9.12	The provisions in these Terms shall apply to the exclusion of any other 
# terms that the Licensee seeks to impose or incorporate, or which are implied by 
# law, trade custom, practice or course of dealing.

# 9.13	Any words following the terms including, include, in particular, for 
# example or any similar expression shall be interpreted as illustrative and 
# shall not limit the sense of the words preceding those terms.

# 9.14	These Terms are drafted in the English language. Any notice or other 
# communication given under or in connection with these Terms shall be in 
# English. All other documents provided under or in connection with these Terms 
# shall be in English, or accompanied by a certified English translation. The 
# English language version of these Terms and any notice, communication or other 
# document relating to these Terms shall prevail over any translation and any 
# version in any other language.



import dimos
import parmed
import torch
import openmm as mm
import openmm.app as mmapp
import numpy as np
from matplotlib import pyplot as plt

# torch.set_default_device("cuda")
torch.set_default_dtype(torch.float64)
prmtop_filename = "./examples/amber/ala2_solv.prmtop"
rst7_filename = "./examples/amber/ala2_solv.rst7"

timestep = 0.5
temperature = 300
initial_temperature = 600

rst7_file = parmed.amber.Rst7(rst7_filename)
positions = rst7_file.positions
positions = torch.tensor(rst7_file.positions.value_in_unit(parmed.unit.angstrom))


system = dimos.ff.AmberForceField(prmtop_filename, nonbonded_type="Cutoff", dispersion_correction=True)
integrator = dimos.integrators.VelocityVerlet(timestep, system)
simulation = dimos.simulation.MDSimulation(system, integrator, positions, temperature=initial_temperature)

openmm_prmtop = mmapp.AmberPrmtopFile(prmtop_filename)
openmm_system = openmm_prmtop.createSystem(
    removeCMMotion=False,
    constraints=None,
    rigidWater=False,
    implicitSolvent=None,
    nonbondedMethod=mmapp.forcefield.CutoffPeriodic,
    nonbondedCutoff=9 * mm.unit.angstrom,
    switchDistance=7.5 * mm.unit.angstrom,
    solventDielectric=78.5)

openmm_integrator = mm.VerletIntegrator(timestep * mm.unit.femtoseconds)
openmm_platform = mm.Platform.getPlatformByName("Reference")

openmm_simulation = mmapp.Simulation(openmm_prmtop.topology, openmm_system, openmm_integrator, platform=openmm_platform)
openmm_simulation.context.setPositions(rst7_file.positions)
openmm_simulation.context.setVelocitiesToTemperature(initial_temperature)
vel = openmm_simulation.context.getState(getVelocities=True).getVelocities(asNumpy=True)
simulation.vel = torch.tensor(vel.value_in_unit(mm.unit.angstrom / mm.unit.femtoseconds)) * dimos.constants.FS_TO_INTERNAL


total_iterations = 100  # 0
measure_every = 10

output_file = open("configs_vv.pdb", "w")

measurements = np.empty((total_iterations, 4))

for t in range(total_iterations):
    print(t)
    simulation.step(measure_every)
    openmm_simulation.step(measure_every)

    state = openmm_simulation.context.getState(getEnergy=True)
    print(t, simulation.sys.calc_energy(simulation.pos, simulation.neighborlist).item(),
          dimos.utils.kinetic_energy(simulation.vel, simulation.sys).item(),
          dimos.utils.measure_temperature(simulation.vel, simulation.sys).item(),
          state.getPotentialEnergy().value_in_unit(mm.unit.kilocalories_per_mole),
          state.getKineticEnergy().value_in_unit(mm.unit.kilocalories_per_mole),
          flush=True)

    measurements[t, 0] = simulation.sys.calc_energy(simulation.pos, simulation.neighborlist).item()
    measurements[t, 1] = dimos.utils.kinetic_energy(simulation.vel, simulation.sys).item()
    measurements[t, 2] = state.getPotentialEnergy().value_in_unit(mm.unit.kilocalories_per_mole)
    measurements[t, 3] = state.getKineticEnergy().value_in_unit(mm.unit.kilocalories_per_mole)

    simulation.sys.write_pdb(simulation.pos, output=output_file, unwrap=True)

s1 = measurements[:, 0] + measurements[:, 1]  # dimos
s2 = measurements[:, 2] + measurements[:, 3]  # OpenMM
t = np.arange(0, total_iterations)

fig, axs = plt.subplots(2, 1, layout='constrained')
axs[0].plot(t, s1, label="dimos")
axs[0].plot(t, s2, label="OpenMM")

axs[0].set_ylabel(r'$E_{\mathrm{tot}}$')
axs[0].grid(True)
axs[0].legend()

axs[1].plot(t, s1 / s2 / (s1[0] / s2[0]))
axs[1].set_ylabel(r"Normalized $E_{\mathrm{tot}}$(dimos)/$E_{\mathrm{tot}}$(OpenMM)")
axs[1].set_xlabel('Iteration')

plt.savefig("energy_conservation.pdf")
