import sys
import itertools
import math

import gsd.hoomd
import hoomd
import numpy as np

import os
import codecs
import h5py

from scipy.spatial.transform import Rotation


cwd = os.getcwd()
os.chdir(cwd)

# Input/output
if len(sys.argv) != 5:
        print("\033[1;31mThe format of the arguments is %s contour_length angle_between_chain_ends number_of_beads betaP\033[0m" % sys.argv[0])

        sys.exit()


lc       = float(sys.argv[1]) #chain contour length
ang      = float(sys.argv[2]) #angle between chain extremities (in degrees)
nb       = int  (sys.argv[3]) #number of beads per chain
betaP    = float(sys.argv[4]) #isobaric Pressure

# nb must be > lc for adjacent beads not to be disconnected
if lc >= nb:
        nb = int(np.ceil(lc))+1

        print("\033[1;36mWarning: lc is too large for the chosen number of beads. Setting nb to %d\033[0m" % nb)

#fz = -1./(nb*lg) #fz is force in z direction (gravity)


# Setup simulation
hoomd.device.CPU(num_cpu_threads=None, communicator=None, message_filename=None, notice_level=2)

# Setup particle geometry
ang_rad        = ang*np.pi/180.

radius         = lc/(np.pi-ang_rad)
gamma          = (np.pi-ang_rad)/2.

angs           = np.linspace(-gamma, gamma, num=nb)
positions      = np.zeros((nb,3))

positions[:,0] = radius*(1.-np.cos(angs))
positions[:,2] = radius*np.sin(angs)

centre_of_mass = positions.mean(axis=0, keepdims=True)
positions     -= centre_of_mass

# Minimal unit cell dimensions to fully encapsulate an individual particle
posx,posy,posz = positions.T

lx  = (posx.max()-posx.min()+1.)
ly  = (posy.max()-posy.min()+1.)
lz  = (posz.max()-posz.min()+1.)

# Particle volume
v   = np.pi*0.5**2*lc + np.pi*0.5**3 *(4/3)

# Inertia tensor and principal frame/moments
inertia      = np.zeros((3,3))

inertia[0,0] = (posy**2+posz**2).sum()
inertia[1,1] = (posx**2+posz**2).sum()
inertia[2,2] = (posx**2+posy**2).sum()

inertia[0,1] = -(posx*posy).sum()
inertia[0,2] = -(posx*posz).sum()
inertia[1,2] = -(posy*posz).sum()

inertia[1,0] = inertia[0,1]
inertia[2,0] = inertia[0,2]
inertia[1,2] = inertia[2,1]

moms,frame   = np.linalg.eigh(inertia)

# Enforce right-handedness of principal frame
if np.linalg.det(frame) < 0: frame[:,0] *= -1.
# Correct for non-ponctual charge distribution for numerical stability
moms[moms.argmin()] = max(1/2.*1.0*0.5**2, moms.min())

# Project positions in principal frame of inertia and convert to quaternion
positions = np.dot(positions, frame)
frame     = Rotation.from_matrix(frame)
quat      = frame.as_quat()

fn = os.path.join(os.getcwd(), 'lattice2.gsd')

fn = os.path.join(os.getcwd(), 'test2.gsd')

fn = os.path.join(os.getcwd(), 'trajectory2.gsd')

fn = os.path.join(os.getcwd(), 'log2.h5')

import fresnel
import IPython

device = fresnel.Device()
tracer = fresnel.tracer.Path(device=device, w=300, h=300)


def render(snapshot):
    
    central_color = fresnel.color.linear([252 / 255, 41 / 255, 0 / 255])
    constituent_color = fresnel.color.linear([93 / 255, 210 / 255, 252 / 255])

    L = snapshot.configuration.box[0]
    scene = fresnel.Scene(device)
    geometry = fresnel.geometry.Sphere(
        scene, N=len(snapshot.particles.position), radius=0.5
    )
    geometry.material = fresnel.material.Material(
        color=[0, 0, 0], roughness=0.5, primitive_color_mix=1.0
    )
    geometry.position[:] = snapshot.particles.position[:]
    geometry.color[snapshot.particles.typeid[:] == 0] = central_color
    geometry.radius[snapshot.particles.typeid[:] == 0] = 0.3
    geometry.color[snapshot.particles.typeid[:] == 1] = constituent_color
    geometry.outline_width = 0.04
    fresnel.geometry.Box(scene, [L, L, L, 0, 0, 0], box_radius=0.02)

    scene.lights = [
        fresnel.light.Light(direction=(0, 0, 1), color=(0.8, 0.8, 0.8), theta=math.pi),
        fresnel.light.Light(
            direction=(1, 1, 1), color=(1.1, 1.1, 1.1), theta=math.pi / 3
        ),
    ]
    scene.camera = fresnel.camera.Orthographic(
        position=(L * 2, L, L * 2), look_at=(0, 0, 0), up=(0, 1, 0), height=L * 1.4 + 1
    )
    scene.background_alpha = 1
    scene.background_color = (1, 1, 1)
    return IPython.display.Image(tracer.sample(scene, samples=500)._repr_png_())

nz = 3
nx = 50
ny = 40
spacingz=lz+2.
spacingx = lx+1.
spacingy= ly+2.
Lz = nz * spacingz
Lx = nx * spacingx
Ly = ny * spacingy
N_particles=nz*nx*ny
x = np.linspace(-Lx/2, Lx/2, nx, endpoint=False)
y = np.linspace(-Ly/2,Ly/2, ny, endpoint=False)
z = np.linspace(-Lz/2, Lz/2, nz, endpoint=False)
position = list(itertools.product(x, y, z))
position = np.array(position) + [spacingx/2, spacingy/2, lz/2]
print('dimensions at start=', Lx, Ly, Lz, ', N_particles=', N_particles, ', nx, ny, nz=', nx, ny, nz)
print('lx ly lz', lx, ly, lz)
print('spacings:', spacingx, spacingy, spacingz)

frame = gsd.hoomd.Frame()
frame.particles.N = N_particles
frame.particles.types = ['R', 'A']
frame.particles.position = position[0:N_particles, :]
frame.particles.typeid = [0] * N_particles
frame.configuration.box = [Lx, Ly, Lz, 0, 0, 0]

frame.particles.mass = [nb] * N_particles
frame.particles.moment_inertia = [moms] * N_particles
frame.particles.orientation = [quat] * N_particles

with gsd.hoomd.open(name='lattice2.gsd', mode='x') as f:
    f.append(frame)

rigid = hoomd.md.constrain.Rigid()

orientations=np.zeros((nb, 4))
orientations[:,0] = 1.
print(orientations)

rigid.body['R'] = {
    'constituent_types': list(('A'*nb)),
    'positions':list(([tuple(row) for row in positions])),
    'orientations': list(([tuple(row) for row in orientations]))
}

simulation = hoomd.Simulation(device=hoomd.device.CPU(), seed=4)
simulation.create_state_from_gsd(filename='lattice2.gsd')
render(simulation.state.get_snapshot())

rigid.create_bodies(simulation.state)

render(simulation.state.get_snapshot())

integrator = hoomd.md.Integrator(dt=0.005, integrate_rotational_dof=True)
integrator.rigid = rigid
simulation.operations.integrator = integrator

simulation.run(0)
render(simulation.state.get_snapshot())

hoomd.write.GSD.write(state=simulation.state, mode='xb', filename='test2.gsd')

cell = hoomd.md.nlist.Cell(buffer=0, exclusions=['body'])
lj = hoomd.md.pair.LJ(nlist=cell)

lj.params[('A', 'A')] = dict( epsilon=5.0, sigma=1.0)
lj.r_cut[('A', 'A')]=2**(1/6.)
lj.params[('R', ['A','R'])] = dict(epsilon=1.0, sigma=1.0)
lj.r_cut[('R', ['A', 'R'])]=0

integrator.forces.append(lj)

p = betaP

rigid_centers_and_free = hoomd.filter.Rigid(('center', 'free'))
npt = hoomd.md.methods.ConstantPressure(
    filter=rigid_centers_and_free, S=betaP, tauS=5.0, couple="none", thermostat=hoomd.md.methods.thermostats.Bussi(kT=1.0, tau=0.5)
)
integrator.methods.append(npt)

simulation.state.thermalize_particle_momenta(filter=rigid_centers_and_free,
                                             kT=1.0)
simulation.run(0)
simulation.run(1000)

thermodynamic_properties = hoomd.md.compute.ThermodynamicQuantities(
    filter=rigid_centers_and_free)

simulation.operations.computes.append(thermodynamic_properties)

logger = hoomd.logging.Logger(categories=['scalar', 'sequence'])
logger.add(obj=thermodynamic_properties, quantities=['potential_energy',
                                                          'translational_kinetic_energy',
                                                          'rotational_kinetic_energy'])
hdf5_writer = hoomd.write.HDF5Log(trigger=hoomd.trigger.Periodic(1000),
                                  filename='log2.h5',
                                  mode='x',
                                  logger=logger)
simulation.operations.writers.append(hdf5_writer)

gsd = hoomd.write.GSD(trigger=hoomd.trigger.Periodic(1_000),
                        filename='trajectory2.gsd')
gsd.maximum_write_buffer_size = 1500000
simulation.operations.writers.append(gsd)

simulation.run(8e8)