# -*- coding: utf-8 -*-
"""Демо: коллективные операции scatter / gather / reduce / bcast.

Запуск из корня курса:
    mpiexec -n 4 python lab08_mpi/examples/03_collectives_mpi.py

Показывает все четыре коллективные операции на маленьком примере.
"""
from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# --- bcast: один -> всем ------------------------------------------------
config = {"task": "summa_kvadratov", "n": 100} if rank == 0 else None
config = comm.bcast(config, root=0)
print(f"ранг {rank}: config = {config}")

# --- scatter: СПИСОК КУСКОВ (ровно size штук), ранг i получает i-й --------
if rank == 0:
    # каждый кусок — свой список; scatter ждёт список из size кусков!
    numbers = [[i * 10 + j for j in range(10)] for i in range(size)]
else:
    numbers = None
my_chunk = comm.scatter(numbers, root=0)    # ранг i получает свой кусок
print(f"ранг {rank}: мой кусок {my_chunk[0]}..{my_chunk[-1]} ({len(my_chunk)} шт.)")

# --- каждый ранг считает СВОЮ часть --------------------------------------
my_sum = sum(x * x for x in my_chunk)
print(f"ранг {rank}: сумма квадратов моего куска = {my_sum}")

# --- gather: все части -> список на ранге 0 ------------------------------
partial_sums = comm.gather(my_sum, root=0)
if rank == 0:
    print(f"ранг 0: частичные суммы = {partial_sums}")
    print(f"ранг 0: итого через sum(gather) = {sum(partial_sums)}")

# --- reduce: то же самое, но сразу со свёрткой (проще и быстрее) ----------
total = comm.reduce(my_sum, op=MPI.SUM, root=0)
if rank == 0:
    print(f"ранг 0: итого через reduce(MPI.SUM) = {total}")
    print("\ngather+sum и reduce(SUM) дали одно и то же — reduce короче.")
    print("В задании mpi_pi вы используете reduce для суммы попаданий.")