# -*- coding: utf-8 -*-
"""Демо: что можно, а что нельзя передать в Pool (pickle).

Запуск из корня курса:
    python lab04_multiprocessing/examples/02_pickle_demo.py

Вложенные функции и лямбды не пиккелятся — PicklingError.
Модульные функции — можно. Отсюда правило: kernel для Pool живёт
на уровне модуля (как dispatch_kernel в common.tasks).
"""
import multiprocessing
import os
import sys

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

from common.tasks import dispatch_kernel, smoke_params


def module_level_double(x):          # МОЖНО: функция уровня модуля
    return x * 2


def make_nested():
    def nested_double(x):            # НЕЛЬЗЯ: вложенная функция
        return x * 2
    return nested_double


def main():
    ctx = multiprocessing.get_context()

    print("--- модульная функция: работает ---")
    with ctx.Pool(2) as pool:
        print(f"  pool.map(module_level_double, [1,2,3]) = "
              f"{pool.map(module_level_double, [1, 2, 3])}")

    print("\n--- lambda: PicklingError ---")
    try:
        with ctx.Pool(2) as pool:
            pool.map(lambda x: x * 2, [1, 2, 3])
    except Exception as e:
        print(f"  {type(e).__name__}: {str(e)[:90]}...")

    print("\n--- kernel задачи через dispatch_kernel: работает ---")
    task_name = "primes"
    chunks = None
    task = __import__("common.tasks", fromlist=["TASKS"]).TASKS[task_name]
    data = task["build"](smoke_params(task_name))
    chunks = task["split"](data, 2)
    payload = [(task_name, ch) for ch in chunks]
    with ctx.Pool(2) as pool:
        partials = pool.map(dispatch_kernel, payload)
    print(f"  подсчёт простых ≤ 2000: {task['combine'](partials)} "
          f"(ожидалось 303)")

    print("\nВывод: всё, что уходит воркерам, должно пиккелиться —")
    print("поэтому в этом курсе kernel диспетчеризуется по имени задачи.")


if __name__ == "__main__":
    main()