import hoomd
import hoomd.md as md
import hoomd.polymd as polymd
import numpy as np
from numpy.random import uniform, seed
import gsd.hoomd
import os
import datetime

seed(0)
simulation = hoomd.Simulation(device = hoomd.device.CPU(notice_level = 2), seed = 0)

# Set up the parameter for molecular dynamics
deltat = 0.002
totalsteps = 10/deltat
snapshots = 1000
kT = 0.25

# Set up the filename and log period
temp_gsd_filename = 'MDrun_temp.gsd'
result_gsd_filename = 'MDrun.gsd'
log_filename = 'MDrun.log'
log_print_period = 1000

# Initialize Our Own Configuration using a Snapshot
rho = 1.00
LParticles = 32
NParticles = LParticles**2
Length = (NParticles / rho) ** 0.5
dmax = 1.0
dmin = 0.5

def placePolydisperseOnSquare(position, diameter):
    for i in range(LParticles):
        for j in range(LParticles):
            position[i*LParticles+j,0] = Length*(i/LParticles-0.5)
            position[i*LParticles+j,1] = Length*(j/LParticles-0.5)
            position[i*LParticles+j,2] = 0
            diameter[i*LParticles+j] = uniform(dmin,dmax)

frame = gsd.hoomd.Frame()
frame.particles.N = NParticles
position = np.zeros((NParticles, 3))
diameter = np.zeros((NParticles,))
placePolydisperseOnSquare(position, diameter)
frame.particles.position = position
frame.particles.typeid = [0] * NParticles
frame.particles.types = ['A']
frame.configuration.box = [Length, Length, 0, 0, 0, 0]
# Because HOOMD v4 or higher does not provide diameter in evaluating potential, we use charge instead of diameter!
frame.particles.charge = diameter
simulation.create_state_from_snapshot(frame)

# Define integrator and neighboring list
integrator = md.Integrator(dt = deltat)
nl = md.nlist.Cell(buffer = 0.4)

# Set up the pair potential
poly12 = polymd.pair.PolydispersePair(nlist = nl, default_r_cut = 4.0)
poly12.params[('A', 'A')] = dict(v0 = 1.0, eps = 0.2, scaledr_cut = 1.25)
integrator.forces.append(poly12)

# Compute thermodynamic properties
thermodynamic_properties = md.compute.ThermodynamicQuantities(filter = hoomd.filter.All())
simulation.operations.computes.append(thermodynamic_properties)

# Set the Logger and gsd writer
samplingtime = int(totalsteps / snapshots)
if samplingtime == 0:
    samplingtime = 1
logger = hoomd.logging.Logger(categories=['scalar', 'string'])
logger.add(simulation, quantities=['timestep'])
logger.add(thermodynamic_properties, quantities=['kinetic_temperature', 'kinetic_energy', 'potential_energy'])
trigger = hoomd.trigger.Periodic(period = samplingtime, phase = 0)
gsd_writer = hoomd.write.GSD(trigger = trigger, filename = temp_gsd_filename, filter = hoomd.filter.All(), mode='wb', 
                            dynamic=['particles/position', 'particles/charge', 'particles/velocity'])
simulation.operations.writers.append(gsd_writer)

# Show the status of simulation
class Status:
    def __init__(self, simulation):
        self.simulation = simulation

    @property
    def seconds_remaining(self):
        try:
            return (
                self.simulation.final_timestep - self.simulation.timestep
            ) / self.simulation.tps
        except ZeroDivisionError:
            return 0

    @property
    def etr(self):
        return str(datetime.timedelta(seconds=self.seconds_remaining)).split('.',2)[0]

status_logger = hoomd.logging.Logger(categories=['scalar', 'string'])
status_logger.add(simulation, quantities=['timestep', 'tps'])
status = Status(simulation)
status_logger[('Status', 'etr')] = (status, 'etr', 'string')
table_status = hoomd.write.Table(trigger=hoomd.trigger.Periodic(period=log_print_period), logger=status_logger)
simulation.operations.writers.append(table_status)

# Set up NVT thermostat
nvt = md.methods.ConstantVolume(filter = hoomd.filter.All(), thermostat=md.methods.thermostats.MTTK(kT = kT, tau = 50*deltat))
integrator.methods.append(nvt)
simulation.operations.integrator = integrator

# Run simulation
with open(log_filename, 'w') as f:
    table = hoomd.write.Table(trigger = trigger, logger = logger, output = f)
    simulation.operations.writers.append(table)
    simulation.run(totalsteps, write_at_start=True)
    gsd_writer.flush()

# Convert charge to diameter in GSD file
data = gsd.hoomd.open(temp_gsd_filename, 'r')
NParticles = data[0].particles.N
with gsd.hoomd.open(result_gsd_filename, 'w') as f:
    nframe = len(data)
    newframe = gsd.hoomd.Frame()
    for i in range(0, nframe):
        frame = data[i]
        newframe = frame
        newframe.particles.diameter = frame.particles.charge
        newframe.particles.charge = np.zeros((NParticles,))
        f.append(newframe)
os.remove(temp_gsd_filename)