"""
案例:
    演示 张量的 索引操作.

分类:
    简单行列索引
    列表索引
    范围索引
    布尔索引
    多维索引

掌握:
    简单行列索引, 范围索引, 多维索引.
"""
# 导包
import torch

# 1. 设置随机种子.
torch.manual_seed(24)

# 2. 创建随机张量.
t1 = torch.randint(1, 10, (5, 5))
print(f't1: {t1}')
print('-' * 30)

# 3. 演示张量的索引操作.
# 场景1: 简单行列索引, 格式: 张量对象[行, 列]
# 需求1: 获取第2行的数据.
print(t1[1])
print(t1[1, :])  # 效果同上, 这里的:表示 所有列

# 需求2: 所有行的第3列数据
print(t1[:, 2])
print('-' * 30)

# 场景2: 列表索引, 前边的表示行, 后边的表示列.
# 需求1: 返回(0, 1), (1, 2)两个位置的元素.
print(t1[[0, 1], [1, 2]])

# 需求2: 返回(1, 2), (3, 4)两个位置的元素.
print(t1[[1, 3], [2, 4]])

# 需求3: 获取第0, 1行的 1, 2列共4个元素.
print(t1[[[0], [1]], [1, 2]])
print('-' * 30)

# 场景3: 范围索引
# 需求1: 前3行, 前2列.
print(t1[:3, :2])

# 需求2: 第2行到最后一行, 前2列的数据.
print(t1[1:, :2])

# 需求3: 所有奇数行, 偶数列.
print(t1[1::2, ::2])    #1::2：从第1行开始，每隔2行取一行（即取第1、3、5…行）。   ::2：对每行，从第0列开始，每隔2列取一列（即取第0、2、4…列）。
print('-' * 30)

# 场景4: 布尔索引
# print(t1[torch.tensor([True, False, False, True, True]), :])    # 演示布尔写法, 看看就好.

# 需求1: 第3列 大于5的行数据.
print(t1[t1[:, 2] > 5])

# 需求2: 第2行大于5的 列数据.
# 理解1: 在第2行基础上 找到该行中大于5的列索引, 然后找所有行中对应列的元素.
print(t1[:, t1[1, :] > 5])
print(t1[:, t1[1] > 5])     # 效果同上.

# 理解2: 在第2行的基础上, 找该行所有列中大于5的元素
print(t1[1, t1[1, :] > 5])
print('-' * 30)


# 场景5: 多维索引
# 创建3维张量, 即: 2个3行4列的矩阵.
t2 = torch.randint(1, 10, (2, 3, 4))
print(f't2: {t2}')

# 需求1: 获取0轴上的第1个数据.
print(t2[0, :, :])

# 需求2: 获取1轴上的第1个数据.
print(t2[:, 0, :])

# 需求3: 获取2轴上的第1个数据.
print(t2[:, :, 0])

print(t2[:,2, 0])

print(t2[:, :, 2])