// 创建变量（支持梯度计算）
Variable x = new Variable(new NdArray(2.0f)).setName("x");
Variable y = new Variable(new NdArray(3.0f)).setName("y");

// 构建计算表达式 z = (x + y) * x = (2 + 3) * 2 = 10
Variable z = x.add(y).mul(x);

// 自动微分：计算 dz/dx 和 dz/dy
z.backward();

// 查看梯度结果
System.out.println("z的值: " + z.getValue().getNumber()); // 输出: 10.0
System.out.println("x的梯度 dz/dx: " + x.getGrad().getNumber()); // 输出: 5.0 (y + x)
System.out.println("y的梯度 dz/dy: " + y.getGrad().getNumber()); // 输出: 2.0 (x)

// 网络参数设置
int batchSize = 32;
int inputSize = 2;   // 输入维度
int hiddenSize = 10; // 隐藏层大小
int outputSize = 3;  // 输出类别数

// 创建多层感知机：input -> hidden -> output
MlpBlock mlpBlock = new MlpBlock("MLP", batchSize, null,
        inputSize, hiddenSize, outputSize);
Model model = new Model("ClassificationModel", mlpBlock);

// 创建随机输入数据
Variable input = new Variable(
        NdArray.likeRandom(-1, 1, new Shape(batchSize, inputSize))
);

// 前向传播
Variable output = model.forward(input);
System.out.println("输出形状: " + output.getValue().getShape()); // [32, 3]

// 训练参数
int maxEpoch = 100;
int batchSize = 10;
float learningRate = 0.01f;
int threadCount = 4; // 并行线程数

// 创建数据集
SpiralDateSet dataSet = new SpiralDateSet(batchSize);

// 创建模型
MlpBlock block = new MlpBlock("ParallelMLP", batchSize, null, 2, 30, 3);
Model model = new Model("ParallelClassifier", block);

// 配置组件
Optimizer optimizer = new Adam(model, learningRate);
Loss lossFunc = new SoftmaxCrossEntropy();
Monitor monitor = new Monitor();
AccuracyEval evaluator = new AccuracyEval(new Classify(), model, dataSet);

// 创建并行训练器
Trainer trainer = new Trainer(maxEpoch, monitor, evaluator, true, threadCount);
trainer.init(dataSet, model, lossFunc, optimizer);

// 开始并行训练
trainer.parallelTrain(true); // true表示打乱数据

// 创建MoE-GPT模型
MoEGPTModel model = MoEGPTModel.createSmallModel("demo_moe_gpt", vocabSize);

// 模型配置
// Small模型：256维，6层，4个专家，Top-2选择
model.printModelInfo();

// 前向传播
Variable output = model.forward(inputTokens);

// 获取负载均衡损失
float balancingLoss = model.computeTotalLoadBalancingLoss();

// 查看专家使用统计
model.printAllExpertStatistics();

// 高级配置
// Tiny模型：128维，4层，2专家，Top-1
// Medium模型：512维，8层，8专家，Top-2

// 准备语料库
List<String> corpus = Arrays.asList(
    "机器", "学习", "是", "人工", "智能", "的", "重要", "分支",
    "深度", "学习", "是", "机器", "学习", "的", "子", "领域"
);

// 创建Word2Vec模型
Word2Vec word2vec = new Word2Vec(
    "word2vec_model",
    50,     // 词汇表大小
    10,     // 词向量维度
    Word2Vec.TrainingMode.SKIP_GRAM, // Skip-gram模式
    2,      // 上下文窗口大小
    true,   // 使用负采样
    5       // 负样本数量
);

// 构建词汇表和生成训练数据
word2vec.buildVocab(corpus);
List<Word2Vec.TrainingSample> samples = word2vec.generateTrainingSamples(corpus);

// 训练模型
Model model = new Model("word2vec_model", word2vec);
Optimizer optimizer = new SGD(model, 0.01f);
SoftmaxCrossEntropy lossFunc = new SoftmaxCrossEntropy();

// 简化训练循环
for (int epoch = 0; epoch < 100; epoch++) {
    for (Word2Vec.TrainingSample sample : samples) {
        Variable input = new Variable(new NdArray(new float[][]{{sample.input}}));
        Variable target = new Variable(new NdArray(new float[][]{{sample.target}}));
        
        Variable output = model.forward(input);
        Variable loss = lossFunc.loss(target, output);
        
        model.clearGrads();
        loss.backward();
        optimizer.update();
    }
}

// 查找相似词
List<String> similarWords = word2vec.mostSimilar("学习", 3);
System.out.println("与'学习'相似的词: " + similarWords);

// GPT-2参数设置
int vocabSize = 1000;   // 词汇表大小
int dModel = 128;       // 模型维度
int numLayers = 4;      // Transformer层数
int numHeads = 4;       // 注意力头数
int maxSeqLength = 64;  // 最大序列长度

// 创建GPT-2模型
GPT2Model gpt2 = new GPT2Model(
    "gpt2_small",
    vocabSize,
    dModel,
    numLayers,
    numHeads,
    maxSeqLength,
    0.1  // dropout率
);

// 模型初始化
gpt2.init();
gpt2.printModelInfo();

// 生成文本示例
NdArray inputTokens = new NdArray(new float[][]{{1, 2, 3, 4, 5}}); // token IDs
Variable logits = gpt2.forward(new Variable(inputTokens));

// 预测下一个token
int nextToken = gpt2.predictNextToken(inputTokens);
System.out.println("预测的下一个token ID: " + nextToken);

// 创建环境
Environment env = new CartPoleEnvironment();

// 创建DQN智能体
DQNAgent agent = new DQNAgent(
    "CartPole_DQN",
    env.getStateDim(),      // 状态维度：4
    env.getActionDim(),     // 动作维度：2
    new int[]{128, 128},    // 隐藏层尺寸
    0.001f,                 // 学习率
    1.0f,                   // 初始探索率
    0.99f,                  // 折扣因子
    32,                     // 批次大小
    10000,                  // 缓冲区大小
    100                     // 目标网络更新频率
);

// 训练循环
for (int episode = 0; episode < 1000; episode++) {
    Variable state = env.reset();
    int totalReward = 0;
    
    while (!env.isDone()) {
        Variable action = agent.selectAction(state);
        Environment.StepResult result = env.step(action);
        
        Experience experience = new Experience(
            state, action, result.getReward(), 
            result.getNextState(), result.isDone()
        );
        
        agent.learn(experience);
        state = result.getNextState();
        totalReward += result.getReward().getValue().getNumber();
    }
    
    if (episode % 100 == 0) {
        System.out.println("Episode " + episode + ": Reward=" + totalReward + 
                          ", Epsilon=" + String.format("%.3f", agent.getCurrentEpsilon()));
    }
}

// 创建GridWorld环境
Environment env = GridWorldEnvironment.createSimpleMaze();

// 创建REINFORCE智能体
REINFORCEAgent agent = new REINFORCEAgent(
    "GridWorld_REINFORCE",
    env.getStateDim(),      // 状态维度：2
    env.getActionDim(),     // 动作维度：4
    new int[]{64, 64},      // 隐藏层尺寸
    0.01f,                  // 学习率
    0.99f,                  // 折扣因子
    true                    // 使用基线
);

// 训练循环
for (int episode = 0; episode < 2000; episode++) {
    Variable state = env.reset();
    float episodeReward = 0;
    
    while (!env.isDone()) {
        Variable action = agent.selectAction(state);
        Environment.StepResult result = env.step(action);
        
        Experience experience = new Experience(
            state, action, result.getReward(), 
            result.getNextState(), result.isDone()
        );
        
        agent.learn(experience); // 存储经验
        state = result.getNextState();
        episodeReward += result.getReward().getValue().getNumber();
    }
    
    // REINFORCE在回合结束时学习
    agent.learnFromEpisode();
    
    if (episode % 200 == 0) {
        System.out.println("Episode " + episode + ": Reward=" + episodeReward);
    }
}

// 保存模型
model.saveModel("models/my_model.model");
model.saveModelCompressed("models/my_model_compressed.model");
model.saveParameters("models/parameters.params");

// 保存训练检查点
model.saveCheckpoint("checkpoints/epoch_100.ckpt", 100, 0.025f);

// 加载模型
Model loadedModel = Model.loadModel("models/my_model.model");

// 从检查点恢复
Model restoredModel = Model.resumeFromCheckpoint("checkpoints/epoch_100.ckpt");

// 模型信息管理
model.setDescription("这是一个图像分类模型");
model.updateTrainingInfo(100, 0.025f, "Adam", 0.001f);
model.addMetric("accuracy", 0.95f);
model.printModelInfo();

// 导出JSON报告
ModelInfoExporter.exportToJson(model, "reports/model_info.json");

// 创建带批量归一化的SimpleConvNet
SimpleConvNet.ConvNetConfig config = new SimpleConvNet.ConvNetConfig()
    .filterNums(32, 64, 128)
    .useBatchNorm(true)
    .useResidual(true)
    .dropoutRate(0.3f);

SimpleConvNet convNet = SimpleConvNet.buildCustomConvNet(
    "enhanced_cnn", 3, 32, 32, 10, config);

// 使用深度可分离卷积
DepthwiseSeparableConvLayer dsConv = new DepthwiseSeparableConvLayer(
    "ds_conv", new Shape(16, 64, 32, 32), 128, 3, 1, 1);

// 批量归一化层
BatchNormLayer batchNorm = new BatchNormLayer("bn", 64, true);
model.saveCheckpoint("checkpoints/epoch_100.ckpt", 100, 0.025f);

// 加载模型
Model loadedModel = Model.loadModel("models/my_model.model");

// 从检查点恢复
Model restoredModel = Model.resumeFromCheckpoint("checkpoints/epoch_100.ckpt");

// 模型信息管理
model.setDescription("这是一个图像分类模型");
model.updateTrainingInfo(100, 0.025f, "Adam", 0.001f);
model.addMetric("accuracy", 0.95f);
model.printModelInfo();

// 导出JSON报告
ModelInfoExporter.exportToJson(model, "reports/model_info.json");