using Flux

# Flux 的核心理念：
# 1. 模型就是函数组合
# 2. 损失函数就是普通 Julia 函数
# 3. 训练循环就是 Julia 循环
# 4. 自动微分通过 Zygote.jl 实现

# 最简单的模型
model = Dense(10 => 5)      # 全连接层：10维输入 → 5维输出
x = rand(Float32, 10)       # 输入
y = model(x)                # 前向传播，5维输出

# 查看模型参数
ps = Flux.params(model)
length(ps)                   # 参数组数

using Flux

# Chain: 将多个层串联成一个模型
model = Chain(
    Dense(784 => 256, relu),    # 输入层 → 隐藏层1
    Dense(256 => 128, relu),    # 隐藏层1 → 隐藏层2
    Dense(128 => 10),           # 隐藏层2 → 输出层
    softmax                     # 概率输出
)

# 前向传播
x = rand(Float32, 784)
y = model(x)                   # 10维概率向量

# 批量输入（矩阵的每一列是一个样本）
X = rand(Float32, 784, 32)     # 32 个样本
Y = model(X)                   # 10×32 矩阵

using Flux

# 分类：交叉熵损失
logits = randn(Float32, 10, 32)    # 10 类，32 个样本
labels = Flux.onehotbatch(rand(0:9, 32), 0:9)  # one-hot 标签
loss = Flux.logitcrossentropy(logits, labels)

# 二分类：二元交叉熵
y_pred = sigmoid.(randn(Float32, 1, 32))
y_true = Float32.(rand(Bool, 1, 32))
loss_binary = Flux.binarycrossentropy(y_pred, y_true)

# 回归：均方误差
y_pred = randn(Float32, 1, 32)
y_true = randn(Float32, 1, 32)
loss_mse = Flux.mse(y_pred, y_true)

# 自定义损失函数
function my_loss(model, x, y)
    ŷ = model(x)
    return Flux.mse(ŷ, y) + 0.01 * sum(abs, Flux.params(model))  # L1 正则
end

using Flux

# 创建模型和数据
model = Chain(Dense(10 => 5, relu), Dense(5 => 1))
X = randn(Float32, 10, 100)
Y = randn(Float32, 1, 100)

# 定义损失
loss(m, x, y) = Flux.mse(m(x), y)

# 优化器
opt = Adam(0.001)          # Adam（最常用）
# opt = SGD(0.01)          # SGD
# opt = AdaGrad(0.01)      # AdaGrad
# opt = RMSProp(0.001)     # RMSProp
# opt = AdaDelta(0.9)      # AdaDelta
# opt = NAdam()            # NAdam

# 获取参数
ps = Flux.params(model)

# 训练循环
data = [(X, Y)]            # 数据必须是 (输入, 标签) 的元组数组
for epoch in 1:1000
    Flux.train!(loss, ps, data, opt)
    if epoch % 100 == 0
        println("Epoch $epoch: loss = $(loss(model, X, Y))")
    end
end

using Flux, Zygote

opt_state = Flux.setup(Adam(0.001), model)

for epoch in 1:1000
    grads = gradient(model) do m
        loss(m, X, Y)
    end
    Flux.update!(opt_state, model, grads[1])
end

using Flux

# 方法1：使用 @layer 宏
struct MyLayer
    W::AbstractMatrix
    b::AbstractVector
end

function MyLayer(in::Int, out::Int)
    W = randn(Float32, out, in) .* sqrt(2.0 / in)
    b = zeros(Float32, out)
    return MyLayer(W, b)
end

function (layer::MyLayer)(x)
    return layer.W * x .+ layer.b
end

Flux.@layer MyLayer   # 注册为 Flux 层（自动支持 GPU 和参数收集）

# 方法2：使用 Chain 和闭包
custom_activation(x) = x * tanh(x)     # 自定义激活函数
model = Chain(Dense(10 => 5, custom_activation), Dense(5 => 1))

using Flux, CUDA

# 检查 GPU 是否可用
CUDA.functional()    # true

# 将模型移到 GPU
model = Chain(Dense(784 => 256, relu), Dense(256 => 10))
gpu_model = gpu(model)      # 移到 GPU

# 将数据移到 GPU
X_gpu = gpu(rand(Float32, 784, 32))
Y_gpu = gpu(Flux.onehotbatch(rand(0:9, 32), 0:9))

# GPU 上训练
Y_pred = gpu_model(X_gpu)

# 移回 CPU
cpu_model = cpu(gpu_model)
X_cpu = cpu(X_gpu)

# 自动检测设备
device = Flux.gpu    # 如果有 GPU 用 gpu，否则自动用 cpu
model = device(model)

using Flux

# 创建数据
X = randn(Float32, 10, 1000)
Y = Flux.onehotbatch(rand(0:9, 1000), 0:9)

# DataLoader: 自动分批、打乱
loader = Flux.DataLoader((X, Y);
    batchsize = 32,
    shuffle = true,      # 每个 epoch 打乱
    partial = true        # 允许最后一批不足 batchsize
)

# 训练循环
model = Chain(Dense(10 => 5, relu), Dense(5 => 10), softmax)
opt = Adam(0.001)
loss(m, x, y) = Flux.logitcrossentropy(m(x), y)

for epoch in 1:10
    for (x_batch, y_batch) in loader
        grads = gradient(model) do m
            loss(m, x_batch, y_batch)
        end
        Flux.update!(opt, model, grads[1])
    end
end

using Flux

# CNN 模型（处理 28×28 灰度图像）
cnn_model = Chain(
    # 卷积块 1
    Conv((3, 3), 1 => 16, relu, pad=SamePad()),
    MaxPool((2, 2)),
    BatchNorm(16),

    # 卷积块 2
    Conv((3, 3), 16 => 32, relu, pad=SamePad()),
    MaxPool((2, 2)),
    BatchNorm(32),

    # 卷积块 3
    Conv((3, 3), 32 => 64, relu, pad=SamePad()),
    MaxPool((2, 2)),
    BatchNorm(64),

    # 全连接层
    Flux.flatten,
    Dense(64 * 3 * 3 => 128, relu),
    Dropout(0.5),
    Dense(128 => 10)
)

# 输入：28×28×1×batch_size
X = rand(Float32, 28, 28, 1, 32)    # 32 张图片
Y_pred = cnn_model(X)                # 10×32 矩阵

# 模型参数统计
total_params = sum(length, Flux.params(cnn_model))
println("总参数量: $total_params")

using Flux

# LSTM 序列模型
model = Chain(
    LSTM(10 => 32),       # 输入 10 维，隐藏 32 维
    Dense(32 => 5)        # 输出 5 类
)

# 序列输入：features × timesteps × batch_size
X = rand(Float32, 10, 20, 16)   # 10 维特征，20 个时间步，16 个样本

# 注意：LSTM 需要逐时间步处理
Flux.reset!(model)        # 重置隐藏状态

# 方式1：取最后时间步
for t in 1:20
    y = model(X[:, t, :])
end
output = X[:, end, :]     # 最终输出

# 方式2：使用 Recur 包装（自动维护状态）
lstm = Flux.Recur(LSTM(10 => 32))
for t in 1:20
    h = lstm(X[:, t, :])
end

using Flux, BSON

# 方法1：保存参数（推荐）
model = Chain(Dense(10 => 5, relu), Dense(5 => 1))
ps = Flux.state(model)    # 获取模型状态
BSON.@save "model.bson" ps

# 加载参数
BSON.@load "model.bson" ps
Flux.loadmodel!(model, ps)

# 方法2：使用 JLD2
using JLD2
@save "model.jld2" model
@load "model.jld2" model

# 方法3：保存为纯 Julia 格式
using Serialization
open("model.jls", "w") do io
    Serialization.serialize(io, model)
end
model_loaded = open(deserialize, "model.jls")

using Flux, MLDatasets, Statistics

# 加载 MNIST 数据集
train_x, train_y = MLDatasets.MNIST(split=:train)[:]
test_x, test_y = MLDatasets.MNIST(split=:test)[:]

# 数据预处理
train_x = Float32.(reshape(train_x, 28, 28, 1, :))   # 28×28×1×60000
test_x = Float32.(reshape(test_x, 28, 28, 1, :))

train_y = Flux.onehotbatch(train_y, 0:9)
test_y = Flux.onehotbatch(test_y, 0:9)

# 构建 CNN
model = Chain(
    Conv((3, 3), 1 => 16, relu, pad=SamePad()),
    MaxPool((2, 2)),
    Conv((3, 3), 16 => 32, relu, pad=SamePad()),
    MaxPool((2, 2)),
    Flux.flatten,
    Dense(32 * 7 * 7 => 128, relu),
    Dropout(0.5),
    Dense(128 => 10)
)

# 损失函数
loss(m, x, y) = Flux.logitcrossentropy(m(x), y)

# 准确率
accuracy(m, x, y) = mean(Flux.onecold(m(x), 0:9) .== Flux.onecold(y, 0:9))

# 数据加载器
train_loader = Flux.DataLoader((train_x, train_y);
    batchsize=128, shuffle=true)

# 优化器
opt_state = Flux.setup(Adam(0.001), model)

# 训练
for epoch in 1:10
    for (x_batch, y_batch) in train_loader
        grads = gradient(model) do m
            loss(m, x_batch, y_batch)
        end
        Flux.update!(opt_state, model, grads[1])
    end

    train_acc = accuracy(model, train_x, train_y)
    test_acc = accuracy(model, test_x, test_y)
    println("Epoch $epoch: Train Acc = $(round(train_acc*100, digits=2))%, Test Acc = $(round(test_acc*100, digits=2))%")
end

# 保存模型
using BSON
ps = Flux.state(model)
BSON.@save "mnist_cnn.bson" ps