import ald
import numpy as np
import h5py

U0 = 1.0
tauR = 1.0
alpha = 1.2
particle = ald.Pareto(U0=U0, tauR=tauR, alpha=alpha)
# sedimentation velocity
Us = 0.1
# sedimentation is downward with a negative sign.
flow = ald.ConstantUy(U=-Us)
Lx = 1.0
# length in y is important
Ly = 20 * U0 * tauR
# wall at y=0.
domain = ald.Box(left=-Lx / 2, right=Lx / 2, bottom=0, top=Ly)


ic = ald.InitialConfig(
    x=ald.Uniform(domain.left, domain.right),
    y=ald.Uniform(domain.bottom, domain.top),
    theta=ald.Uniform(0, 2 * np.pi),
)


cfg = ald.Config(particle, domain, N=204_800, dt=1e-4, Nt=4_000_000)

kernel = ald.RTPAboveWallKernel()
compiler = ald.RTPCompiler(kernel, cfg, flow, ic)


compiler.compile()

simulator = ald.RTPSimulator(cfg, compiler)
if isinstance(particle, ald.Pareto):
    tag = "pareto"
elif isinstance(particle, ald.ExponentialRTP):
    tag = "rtp"
else:
    tag = ""

file = "U{:.3f}tauR{:.3f}sediment_{}.h5".format(U0, tauR, tag)
# create an empty file
# with h5py.File(file, "w") as f:
#    pass

# save configuraton.
# callback controller
# x is periodic, unwrap to save the absolute position
# y no need to do so
configsaver = ald.ConfigSaver(
    ald.RangedRunner.from_backward_count(stop=cfg.Nt, freq=10000, count=200),
    file,
    variables=["x", "y"],
    unwrap=[True, False],
)


# range to compute stats on configuration and print time.
runner = ald.RangedRunner(start=0, stop=cfg.Nt, freq=10000)
callbacks = [configsaver, ald.ETA(runner)]


simulator.run(cfg, callbacks=callbacks)

# save particle, domain and simulation attributes.
cfg.save2h5(file)
# save mean variance of x
# x.save2h5(file, "x")
# y.save2h5(file, "y")
