# 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 torch
import parmed
import dimos
from dimos.integrators import LangevinDynamics
from dimos.mlip.mace import MaceFoundationalSystem
from dimos.ff import AmberForceField
from dimos.mlmm import MLMMsystem
from dimos.simulation import MDSimulation

timestep = 0.5
temperature = 300
initial_temperature = 600
friction_coefficient = 0.1
compile=True

torch.set_default_device("cuda:0")
torch.set_default_dtype(torch.float32)
prmtop_filename = "./examples/amber/ala2_solv.prmtop"
rst7_filename = "./examples/amber/ala2_solv.rst7"

parameters = parmed.amber.AmberParm(prmtop_filename)
ml_atoms = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22], dtype=torch.int32)

num_atoms = len(parameters.atoms)
atomic_numbers = torch.zeros(num_atoms, dtype=torch.int32)
masses = torch.zeros(num_atoms, dtype=torch.get_default_dtype())
names = []

for idx, atom in enumerate(parameters.atoms):
    #if idx in ml_atoms:
        atomic_numbers[idx] = atom.atomic_number
        names.append(atom.type[0])
        masses[idx] = atom.mass

mm_system = AmberForceField(prmtop_filename, nonbonded_type=None, dispersion_correction=True, excluded_bonded_atoms=ml_atoms)
ml_system = MaceFoundationalSystem(masses, atomic_numbers, periodic=mm_system.periodic, box=mm_system.box)

mlmm_system = MLMMsystem(ml_atoms, ml_system, mm_system)

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

integrator = LangevinDynamics(timestep, temperature, friction_coefficient, mlmm_system)
simulation = MDSimulation(mlmm_system, integrator, positions, temperature=initial_temperature)

if compile:
    step = torch.compile(simulation.step, mode="max-autotune-no-cudagraphs")
else:
    step = simulation.step

file_mlmm = open("test_mlmm.pdb", "w")
for iteration in range(100):
    print(iteration)
    step(10)
    #simulation.step(10)
    simulation.sys.mm_system.write_pdb(simulation.pos, file_mlmm)
