# Path: modred/parallel.py
# def get_hostname():
# def get_node_ID():
# def get_num_nodes():
# def get_num_MPI_workers():
# def get_rank():
# def get_num_procs():
# def is_distributed():
# def is_rank_zero():
# def barrier():
# def print_from_rank_zero(msg, output_channel='stdout'):
# def call_from_rank_zero(func, *args, **kwargs):
# def bcast(vals):
# def call_and_bcast(func, *args, **kwargs):
# def find_assignments(tasks, task_weights=None):
# def check_for_empty_tasks(task_assignments):


import os
import numpy as np
import modred as mr
from modred import parallel
from customvector import CustomVector, CustomVecHandle, inner_product




# Create directory for output files
out_dir = 'tutorial_ex6_out'
if not os.path.isdir(out_dir):
    os.makedirs(out_dir)

# Define snapshot handles
direct_snapshots = [
    CustomVecHandle('%s/direct_snap%d.pkl' % (out_dir, i), scale=np.pi)
    for i in mr.range(10)]
adjoint_snapshots = [
    CustomVecHandle('%s/adjoint_snap%d.pkl' % (out_dir, i), scale=np.pi)
    for i in mr.range(10)]

# Create random snapshot data
nx = 50
ny = 30
nz = 20
x = np.linspace(0, 1, nx)
y = np.logspace(1, 2, ny)
z = np.linspace(0, 1, nz) ** 2
if parallel.is_rank_zero():
    for snap in direct_snapshots + adjoint_snapshots:
        snap.put(CustomVector([x, y, z], np.random.random((nx, ny, nz))))

# Path: modred/parallel.py
# def get_hostname():
# def get_node_ID():
# def get_num_nodes():
# def get_num_MPI_workers():
# def get_rank():
# def get_num_procs():
# def is_distributed():
# def is_rank_zero():
# def barrier():
# def print_from_rank_zero(msg, output_channel='stdout'):
# def call_from_rank_zero(func, *args, **kwargs):
# def bcast(vals):
# def call_and_bcast(func, *args, **kwargs):
# def find_assignments(tasks, task_weights=None):
# def check_for_empty_tasks(task_assignments):


# Path: modred/examples/tutorial_ex6.py
import os
import numpy as np
import modred as mr
from modred import parallel
from customvector import CustomVector, CustomVecHandle, inner_product





# Create directory for output files
out_dir = 'tutorial_ex6_out'
if not os.path.isdir(out_dir):
    os.makedirs(out_dir)

# Define snapshot handles
direct_snapshots = [
    CustomVecHandle('%s/direct_snap%d.pkl' % (out_dir, i), scale=np.pi)
    for i in mr.range(10)]
adjoint_snapshots = [
    CustomVecHandle('%s/adjoint_snap%d.pkl' % (out_dir, i), scale=np.pi)
    for i in mr.range(10)]

# Create random snapshot data
nx = 50
ny = 30
nz = 20
x = np.linspace(0, 1, nx)
y = np.logspace(1, 2, ny)
z = np.linspace(0, 1, nz) ** 2
if parallel.is_rank_zero():
    for snap in direct_snapshots + adjoint_snapshots:
        snap.put(CustomVector([x, y, z], np.random.random((nx, ny, nz))))

parallel.barrier()