import nn

model = nn.model.SequentialModel(input_shape=[-1, 784])
model.add(nn.layer.Dense(128, activation=nn.activation.Tanh()))
model.add(nn.layer.Dense(128, activation=nn.activation.Tanh()))
model.add(nn.layer.Dense(10, activation=nn.activation.Softmax()))

model.setup(nn.loss.Cross_Entropy_With_Softmax(), nn.metric.CategoricalAccuracy())

import nn

model.save("dnn.model")  # 使用save保存
nn.model.Model.load("dnn.model")  # 使用load加载

import nn

data = nn.dataset.MNIST()  # 使用MNIST数据集

import nn.dataset.transforms as transforms

trans = transforms.Shuffle().add(transforms.ImageCls())  # 先对数据集做Shuffle操作，再对数据集进行像素分类处理

import parallel_sgd

job = parallel_sgd.ParallelSGD(model, data, trans)

import parallel_sgd
from parallel_sgd.codec.plain import Plain

nodes = parallel_sgd.parse_worker(worker_cnt=1)
res = job.parallel(nodes, codec=Plain, epoch=2)
print(res)

import parallel_sgd
import nn
import nn.dataset.transforms as transforms

from parallel_sgd.codec.plain import Plain

model = nn.model.SequentialModel(input_shape=[-1, 784])
model.add(nn.layer.Dense(128, activation=nn.activation.Tanh()))
model.add(nn.layer.Dense(128, activation=nn.activation.Tanh()))
model.add(nn.layer.Dense(10, activation=nn.activation.Softmax()))

model.setup(nn.loss.Cross_Entropy_With_Softmax(), nn.metric.CategoricalAccuracy())

data = nn.dataset.MNIST()  # 使用MNIST数据集
trans = transforms.Shuffle().add(transforms.ImageCls())  # 先对数据集做Shuffle操作，再对数据集进行像素分类处理

job = parallel_sgd.ParallelSGD(model, data, trans)
nodes = parallel_sgd.parse_worker(worker_cnt=1)

# 两轮训练后，大约可以给出 'Loss': 0.262, 'accuracy': 0.954 的结果
print(job.parallel(nodes, codec=Plain, epoch=2))

import parallel_sgd
import rpc
import network

nodes = parallel_sgd.parse_worker(worker_cnt=9, ps=True)
req = network.Request()

with req.request(nodes) as com:
  rpc.Reclaimer(com).require_client_log()

from typing import Hashable, SupportsFloat, Type, Union, Dict, Tuple

from network import NodeAssignment
from nn.gradient_descent import ADAMOptimizer
from nn.gradient_descent.interface import IGradientDescent
from nn.optimizer.parallel_sgd import PSGDOptimizer, IOptimizer
from parallel_sgd.batch_sync import SynchronizedSGD
from parallel_sgd.batch_sync.sync.interface import ISyncType
from parallel_sgd.codec.interfaces import Codec
from parallel_sgd.profiles import IIDBlockAssignment
from parallel_sgd.profiles.blockassignment.abstract import AbsBlockAssignment


def parallel(self,
             nodes: NodeAssignment,
             redundancy: int = 1,
             block_size: int = 64,
             epoch: int = 10,
             assignment_type: Type[AbsBlockAssignment] = IIDBlockAssignment,
             sync_type: Type[ISyncType] = SynchronizedSGD,
             op_type: Type[IOptimizer] = PSGDOptimizer,
             gd_type: Type[IGradientDescent] = ADAMOptimizer,
             codec: Union[Dict[int, Type[Codec]], Type[Codec]] = None,
             gd_params: Tuple[object] = (),
             ps_codec: Union[Dict[int, Type[Codec]], Type[Codec], None] = None,
             network_bandwidth: int = 1048576,
             mission_title: str = "P-SGD",
             ssgd_timeout_limit: int = 10000,
             codec_extra_parameters: Dict[Hashable, SupportsFloat] = None) -> Dict[str, SupportsFloat]:
  """
      执行并行化。
  :param ssgd_timeout_limit: Sync-SGD等待超时限制，单位为毫秒，数值为整型，控制直接相连的Worker之间的最大
                          执行速度差距。如果两个Worker没有直接通讯，则不做限制。例如：连在同一个参数服务器
                          下的两个Worker，由于没有直接通信，则不会限制他们之间的执行速度差异。

  :param network_bandwidth: 可用的网络带宽，用作计算预估传输时间，设置 pre_commit 超时计时器。

  :param mission_title:   任务标题，作为本次任务的log文件文件名。

  :param nodes:           由 network 模块提供的 NodeAssignment 接口，指示了当前并行化操作调用的节点数目。
                          参数服务器的节点编号由 utils.constant.Parameter_Server 指定，其余工作节点的id
                          从 0 开始依次递增（为整数型）。

  :param block_size:      节点粒度的 Batch 大小，由 codec 控制具体的更新策略，块大小与批次大小并没有具体
                          对应关系。若 codec 在每个 Block 后给出结果，则实践意义上 Block size 和 Batch size 
                          是等价的，例如：异步随机梯度下降算法。
                          若 codec 总是等待所有 Block 的训练完成后再同步结果，则实践意义上 Batch size 
                          等于 Block size 乘以 Block 总数，例如：Ring AllReduce SGD。

  :param epoch:           训练批次数，由 codec 和 sync type 共同决定 epoch 内的同步策略，当使用参数服务器
                          时，参数服务器也参与到同步状态的维持中。
                          若 codec 不允许异步执行，则所有节点都会在同一时刻结束 epoch，若 codec 或 sync type 
                          允许跨批次执行，则节点会根据自己的计算能立先后结束计算。

  :param assignment_type: 样本分配策略，一般与冗余分配结合使用，需要实现了 profiles.ISetting 接口的类型。
                          初衷是提供冗余的数据集分配策略，现可以提供静态数据量分配。
                              主要指定一个内容：数据集是如何划分到每个节点上的。

  :param sync_type:       同步方式。分同步和异步两种，需要实现了 parallel_sgd.sync.IParallelSGD 接口的类型。
                          同步方式下，每个 worker 在调用 get_weights() 获取权重时才会处理接收数据。
                          异步方式下，每个 Worker 收到数据就处理并更新结果集。
                          具体的数据处理流程和结果集更新策略都由 codec 定义。
                              主要指定一个内容：Worker之间的连接是同步的还是异步的，是否会存在某个worker
                              在运行时突然接收到外部数据。

  :param gd_type:         梯度生成策略，实现了 nn.gradient_descent.IGradientDescent 接口的类型。
                          这里定义了梯度与迭代增量的关系，我们一般不会使用计算出的梯度直接执行反向传播算法，
                          而是对计算出的梯度进行一些处理之后再进行迭代。
                              内置的范式：Adam、RMSProp、AdaDelta、GradientDecay、SGD。
                              主要指定一个内容：如何处理原始梯度，使其更满足模型的需求。

  :param op_type:         梯度处理策略类型，实现了 nn.IOptimizer 接口的类型。
                          这里定义了模型如何对待获取到的更新增量。
                              内置的范式：梯度上升、梯度下降、梯度平均化、并行双重缓冲、参数平均化，联邦学习。
                              主要指定两个内容：1.模型是如何处理更新量\delta w 的，2.模型是如何执行反向传播的。

  :param codec:           编码器类型，实现了 codec.interface.Codec 接口的类型。
                          这个Codec在Worker上运行，用于处理一些编码压缩，稀疏化的操作，这里也定义了
                          Worker要和谁交换数据，如何处理收到的数据。
                              内置的范式：P2P交换网络、中心化交换网络、联邦平均化、动态量化、三值化、随机量化、Adam PS。
                              主要指定两个内容：1.如何处理要发送的梯度，2.发送给谁。

  :param redundancy:      冗余设置，适用于能够处理冗余的 codec 和 block assignment。
                          这个参数会被传递给 @assignment_type，并且要求提供的 @codec 能够处理冗余数据。

  :param gd_params:       梯度生成器参数，传给 @gd_type。

  :param ps_codec:        编码器类型，实现了 codec.interface.Codec 接口的类型。
                          这个Codec在参数服务器上运行，包含了合并全局模型的逻辑。
                              内置的范式：中心化交换网络、联邦平均化、动态量化、三值化、随机量化、Adam PS。
                              主要指定两个内容：1.如何处理全局参数合并，2.如何同步各个Worker的状态。

  :param codec_extra_parameters:用于Codec接口识别的其他参数列表，为字典形式。该字典将会存储在每个Worker的
                          codec.GlobalSettings.__global_parameters 参数中，Codec对象可以使用
                          GlobalSettings.get_params(key: str) -> object 函数获取对应key的值。

  :return: Dict，代表全局执行结果，包含平均准确率和损失，以及在Model中定义的其他评判指标。
  """
  pass