import numpy as np
import matplotlib.pyplot as plt
from ase.calculators import eam
from ase.geometry.analysis import Analysis
from ase.io.lammpsdata import write_lammps_data

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

n, m = 4, 4
CH_bond, CC_bond = 1.09, 1.42 # bond lengths in Angstroms
initial_mag = 1.12  # initial magnetic moment in some units

flake = build.graphene_nanoribbon(n=n, m=m, type='zigzag', saturated=False,
                                  C_H=CH_bond, C_C=CC_bond, vacuum=2*CC_bond,
                                  magnetic=False, initial_mag=initial_mag,
                                  sheet=False, main_element='C', saturate_element='H')

ana = Analysis(flake)
CCBonds = ana.get_bonds('C', 'C', unique=True)
bonds = CCBonds[0]

if True:
    with open('wow.data', 'w') as outfile:
        write_lammps_data(outfile, flake)

# flake.calc = eam

x, y, z = flake.positions.T
n_atoms = len(flake)
indices = np.arange(n_atoms)
sequence = indices / (n_atoms - 1)


fig = plt.figure(figsize=[12, 7.5])
ax1  = fig.add_subplot(1, 1, 1, projection='3d', proj_type = 'ortho')
ax1.scatter(x, y, z, s=400, c=sequence)
hw = np.abs(flake.positions).max()
ax1.set_xlim(-0.05*hw, 1.05*hw)
ax1.set_ylim(-0.05*hw, 1.05*hw)
ax1.set_zlim(-0.05*hw, 1.05*hw)
ax1.set_box_aspect([1,1,1])

plt.show()