from mpi4py.MPI import ANY_SOURCE

# Common MPI communication, rank, size
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# Get an OSCARS SR object with nthreads=1 (meaning 1 thread per mpi execution)
osr = oscars.sr.sr(nthreads=1)

# Set a particle beam with non-zero emittance
osr.set_particle_beam(
    energy_GeV=3,
    x0=[0, 0, -1],
    current=0.500,
    sigma_energy_GeV=0.001*3,
    beta=[1.5, 0.8],
    emittance=[0.9e-9, 0.008e-9]
)

# Must set the start and stop time for calculations
osr.set_ctstartstop(0, 2)

# Phase difference between fields in [rad]
phase = osr.pi()/2.

# Clear any existing fields (just good habit in notebook style) and add an undulator field
osr.clear_bfields()
osr.add_bfield_undulator(bfield=[0, 0.7, 0], period=[0, 0, 0.049], nperiods=21, phase=-phase/2.)
osr.add_bfield_undulator(bfield=[0.7, 0, 0], period=[0, 0, 0.049], nperiods=21, phase=+phase/2.)

# Number of particles per node of rank > 1
particles_per_node = 20

# Observation point for spectrum
observation_point = [0, 0, 30]

# Number of points in the spectrum
npoints = 200

# Energy range for spectrum
range_eV = [130, 180]




# Now the MPI fun
if rank == 0:
    # For rank 0 we calculate the ideal single-particle spectrum
    osr.set_new_particle(particle='ideal')
    spectrum_se = osr.calculate_spectrum(obs=observation_point, energy_range_eV=range_eV, npoints=npoints)

    # Weight for each spectrum in summation
    weight = 1. / size

    # Get a new OSCARS SR object to collect the data in
    osr_sum = oscars.sr.sr()

    # Now wait and collect data from all other processes when it comes in
    for i in range(1, size):
        # Get incoming data
        data = comm.recv(source=ANY_SOURCE)

        # Sum the spectra (internally this is a compensated sum)
        osr_sum.add_to_spectrum(spectrum=data, weight=weight)

    # Plot the single-particle and multi-particle data and save to file
    plot_spectra([spectrum_se, osr_sum.get_spectrum()], ['single-electron', 'multi-electron'], ofile='EPUSpectrum_MPI.png', show=False)



else:
    # If not rank 0, calculate the desired spectrum
    data = osr.calculate_spectrum(obs=observation_point, energy_range_eV=range_eV, npoints=npoints, nparticles=particles_per_node)

    # Send results back to rank 0
    comm.send(data, dest=0)