import random

from particlecloud.stuff import Particle


PARTICLE_FILE_COLUMN_NAMES = [
    "number", "mass", "vx", "vy", "vz", "px", "py", "pz",
]

def make_particles(n=10):
    particles = []
    for i in range(n):
        particles.append(
            Particle(
                mass=random.gauss(mu=3., sigma=0.5),
                velocity=(
                    random.gauss(mu=0., sigma=1.5),
                    random.gauss(mu=0., sigma=1.5),
                    random.gauss(mu=0., sigma=1.5),
                ),
                position=(
                    random.gauss(mu=0., sigma=25),
                    random.gauss(mu=0., sigma=25),
                    random.gauss(mu=0., sigma=25),
                )
            )
        )
    return particles


def write_particles_file(filename, n=10):
    """Write a csv file with particle data
    """
    particles = make_particles(n=n)
    with open(filename, "w") as fp:
        fp.write(", ".join(PARTICLE_FILE_COLUMN_NAMES) + ",\n")
        for i, p in enumerate(particles):
            vals = []
            vals.append(p.mass)
            vals.extend(p.velocity)
            vals.extend(p.position)
            line = f"{i}, " + ", ".join(f"{v:7.5f}" for v in vals) + ",\n"
            fp.write(line)


if __name__ == '__main__':
    write_particles_file("sample_data.csv")