import scm
import scm.pyzacros as pz

# Gas species
CO_g = pz.Species("CO")
O2_g = pz.Species("O2")
CO2_g = pz.Species("CO2", gas_energy=-2.337)

# Surface species
s0 = pz.Species("*")  # Empty adsorption site
CO_s = pz.Species("CO*")
O_s = pz.Species("O*")

# Species List
spl = pz.SpeciesList([CO_g, O2_g, CO2_g, s0, CO_s])
spl.append(O_s)
print(spl)

# Lattice setup
lat = pz.Lattice(lattice_type=pz.Lattice.TRIANGULAR, lattice_constant=1.0, repeat_cell=[10, 3])

print(lat)

lat.plot()

# Clusters
CO_p = pz.Cluster(species=[CO_s], energy=-1.3)
O_p = pz.Cluster(species=[O_s], energy=-2.3)
print(CO_p)

# Cluster Expansion
ce = pz.ClusterExpansion([CO_p, O_p])
print(ce)

# Elementary Reactions
CO_ads = pz.ElementaryReaction(
    initial=[s0, CO_g], final=[CO_s], reversible=False, pre_expon=10.0, label="CO_adsorption"
)

O2_ads = pz.ElementaryReaction(
    initial=[s0, s0, O2_g],
    final=[O_s, O_s],
    neighboring=[(0, 1)],
    reversible=False,
    pre_expon=2.5,
    label="O2_adsorption",
)

CO_oxi = pz.ElementaryReaction(
    initial=[CO_s, O_s],
    final=[s0, s0, CO2_g],
    neighboring=[(0, 1)],
    reversible=False,
    pre_expon=1.0e20,
    label="CO_oxidation",
)

mech = pz.Mechanism([O2_ads, CO_ads, CO_oxi])

print(mech)

# LatticeState setup (initial state)
ist = pz.LatticeState(lat, surface_species=spl.surface_species())
ist.fill_sites_random(site_name="StTp1", species="CO*", coverage=0.1)
ist.fill_sites_random(site_name="StTp1", species="O*", coverage=0.1)

print(ist)

ist.plot()

scm.pyzacros.init()

# Settings:
sett = pz.Settings()
sett.random_seed = 953129
sett.temperature = 500.0
sett.pressure = 1.0
sett.snapshots = ("time", 0.1)
sett.process_statistics = ("time", 0.1)
sett.species_numbers = ("time", 0.1)
sett.event_report = "off"
sett.max_steps = "infinity"
sett.max_time = 1.0

sett.molar_fraction.CO = 0.45
sett.molar_fraction.O2 = 0.55

print(sett)

job = pz.ZacrosJob(settings=sett, lattice=lat, mechanism=[CO_ads, O2_ads, CO_oxi], cluster_expansion=[CO_p, O_p])

print(job)

results = job.run()

if job.ok():
    provided_quantities = results.provided_quantities()
    print("nCO2 =", provided_quantities["CO2"])

    results.plot_molecule_numbers(results.gas_species_names())
    results.plot_lattice_states(results.lattice_states())

    pstat = results.get_process_statistics()
    results.plot_process_statistics(pstat, key="number_of_events")

scm.pyzacros.finish()