from ase.build import bulk
from msspec.utils import hemispherical_cluster, get_atom_index
from ase.visualize import view
from datetime import timedelta as td
import numpy as np
from matplotlib import pyplot as plt

Si = bulk('Si', cubic=True)
planes = range(1,10)
natoms = []
allt = []
for emitter_plane in planes:
    cluster = hemispherical_cluster(Si, diameter=40, emitter_plane=emitter_plane, planes=emitter_plane+1)
    N = len(cluster)
    npaths = np.sum([(N-1)**i for i in range(emitter_plane + 1)])
    t = npaths * 1e-6
    
    natoms.append(len(cluster))
    allt.append(t)
    
    values=np.array([])
    conversion = [31557600, 86400, 3600, 60, 1, 1e-3, 1e-6]
    units = np.array([" years", " days", "H", "M", "s", "ms", "µs"])
    for f in conversion:
        values = np.append(values, t // f)
        t -= values[-1] * f
    
    s = ' '.join(["{:.0f}{}".format(*_) for _ in zip(values[values>0][:3], units[values>0][:3])])
    print("Emitter in plane {:d}".format(emitter_plane))
    print("Number of atoms in the cluster: {:d}".format(len(cluster)))
    print("{:d} scattering paths".format(npaths))
    print("{} to compute".format(s))
    print("*"*80)
    

colibri = 1/60    
one_day = 86400
one_year = 365.25 * 86400
roman_era = 1200 * one_year
jurassic_period = 55e6 * one_year
universe_age = 13.8e9 * one_year
    
fig, ax = plt.subplots()

for _ in [colibri, one_day, roman_era, jurassic_period, universe_age]:
    ax.axhline(_, linestyle="--", color="black", lw=1)

ax.semilogy(natoms, allt, color="red", lw=2)

ax.scatter([natoms[0], natoms[6]], [allt[0],allt[6]], marker="o", color="red", lw=3)


plt.ylim(1e-6, 1e20)
plt.xlabel("Number of atoms for emitter in plane 1 to 9")
plt.ylabel("Computation time (s)")
plt.savefig('my_plot.png', transparent=True)
plt.show()