# -*- coding: utf-8 -*-
"""Лабораторная работа 8: распределённые вычисления (mpi4py).

Заполните функции ниже. ИНТЕРФЕЙСЫ МЕНЯТЬ НЕЛЬЗЯ — по ним работают
автотесты (tests/test_lab08.py; если mpi4py не установлен — тесты
пропускаются, но работа сдаётся с запуском через mpiexec).

Запуск из корня курса:
    python lab08_mpi/solution.py                 (один ранг — отладка)
    mpiexec -n 4 python lab08_mpi/solution.py    (4 ранга — замеры)
"""
import os
import random
import sys
import time

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from common.tasks import TASKS, get_variant

try:
    from mpi4py import MPI
    HAVE_MPI = True
except ImportError:
    HAVE_MPI = False

# TODO: впишите свой номер в журнале (1..20)
VARIANT_NUMBER = 0


def _group_for_ranks(chunks, size):
    """Распределить куски по рангам round-robin: список из size групп.

    scatter ожидает список ровно из size элементов, а естественных кусков
    может быть другое число (например, nqueens всегда даёт N кусков) —
    поэтому группируем.
    """
    # TODO: groups = [[] for _ in range(size)]; groups[i % size].append(ch)
    raise NotImplementedError


def mpi_run(task_name: str, params: dict):
    """CPU-задача на всех рангах: scatter групп -> kernel -> gather -> combine.

    Схема:
      ранг 0: build -> split(data, size) -> _group_for_ranks
      scatter раздаёт ГРУППЫ кусков (по одной на ранг)
      каждый ранг: my_partials = [task["kernel"](ch) for ch in my_group]
      gather собирает списки частичных результатов
      ранг 0: развернуть в плоский список -> combine ОДИН раз -> checksum

    Функции пиккелить не нужно (в отличие от multiprocessing) — только данные.

    Для задачи 'life': данные — (grid, steps), параллелятся ПОКОЛЕНИЯ.
    Ранг 0 держит текущее поле; каждый шаг: scatter полос -> kernel ->
    gather -> combine -> новое поле. Число шагов раздаётся через bcast.
    """
    if not HAVE_MPI:
        raise RuntimeError("mpi4py не установлен (см. методичку, раздел Установка)")
    task = TASKS[task_name]
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size = comm.Get_size()

    if task_name == "life":
        # TODO: цикл по поколениям (см. схему выше); в конце bcast поля
        #  с ранга 0 всем, чтобы checksum был у всех рангов
        raise NotImplementedError

    # TODO: общий случай по схеме выше; вернуть checksum на ранге 0 (None на других)
    raise NotImplementedError


def mpi_pi(N: int) -> float:
    """π методом Монте-Карло на всех рангах (reduce).

    Каждый ранг бросает N//size точек с зерном random.Random(100 + rank).
    reduce(MPI.SUM) сворачивает попадания на ранге 0.
    Вернуть π (на ранге 0 — вычисленное значение, на остальных 0.0).
    """
    # TODO: my_n = N // size; цикл бросков; comm.reduce(inside, op=MPI.SUM, root=0)
    raise NotImplementedError


def main():
    if not HAVE_MPI:
        print("mpi4py не установлен. Инструкция — в методичке, раздел «Установка».")
        print("Кратко (Windows): MS-MPI runtime + pip install mpi4py")
        sys.exit(1)

    v = get_variant(VARIANT_NUMBER)
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size = comm.Get_size()

    if rank == 0:
        print(f"Вариант {v['variant']}: {v['describe']}")
        print(f"Запущено рангов: {size}\n")

    # --- pi через reduce ------------------------------------------------------
    pi = mpi_pi(1_000_000)
    if rank == 0:
        print(f"π ≈ {pi:.6f} (Монте-Карло, 1 000 000 бросков суммарно)")

    # --- основная задача --------------------------------------------------------
    comm.barrier()
    t0 = time.perf_counter()
    checksum = mpi_run(v["task_name"], v["cpu_params"])
    dt = time.perf_counter() - t0

    if rank == 0:
        from common.tasks import run_sequential
        ref, _ = run_sequential(v["task_name"], v["cpu_params"], parts=1)
        ok = checksum == ref or (isinstance(checksum, float) and
                                 abs(checksum - ref) < 1e-6 * max(1, abs(ref)))
        print(f"задача: {dt:.3f} c, checksum совпадает с последовательным: {ok}")
        print("\nЗапустите при разном числе рангов для таблицы T(p):")
        print("  mpiexec -n 1 python lab08_mpi/solution.py")
        print("  mpiexec -n 2 python lab08_mpi/solution.py")
        print("  mpiexec -n 4 python lab08_mpi/solution.py")


if __name__ == "__main__":
    main()