from enum import IntEnum, auto
from pathlib import Path
import sys
import torch
import matplotlib.pyplot as plt
import torch.nn.functional as F

sys.path.append(str(Path(__file__).resolve().parent.parent))
import gpu_utils


class CurveKind(IntEnum):
    """生成假数据时使用的曲线"""
    Polynomials = auto()
    Sine = auto()


class DataSource:
    """用于拟合的随机生成的假数据"""

    x: torch.Tensor
    y: torch.Tensor

    def __init__(self, device: torch.device, curve_kind: CurveKind):
        match curve_kind:
            case CurveKind.Polynomials:
                x = torch.linspace(-1, 1, steps=100).reshape(-1, 1)
                y = -x.pow(3) + 2 * x.pow(2) + 0.2 * torch.rand(x.size())
            case CurveKind.Sine:
                # 正弦在0-2之间变化才不是类似线性的
                x = torch.linspace(0, 2, steps=100).reshape(-1, 1)
                y = x.sin() + 0.2 * torch.rand(x.size())

        self.x = x.to(device)
        self.y = y.to(device)


class Net(torch.nn.Module):
    """继承torch的module用于表示网络"""

    def __init__(self, n_feature, n_hidden, n_output):
        super(Net, self).__init__()  #继承_init_功能
        #定理每层用什么样的形式
        self.hidden1 = torch.nn.Linear(n_feature, n_hidden)  #隐藏层线性输出
        self.hidden2 = torch.nn.Linear(n_hidden, n_hidden)  #输出层线性输出
        self.hidden3 = torch.nn.Linear(n_hidden, n_hidden)  #输出层线性输出
        self.predict = torch.nn.Linear(n_hidden, n_output)  #输出层线性输出

    def forward(self, x):  #这同时也是module中的forward功能
        #正向传播输入值，神经网络分析出输出值
        x = F.relu(self.hidden1(x))  #激励函数（隐藏层的线性值）
        x = F.relu(self.hidden2(x))
        x = F.relu(self.hidden3(x))
        x = self.predict(x)  #输出值
        return x


def main():
    device = gpu_utils.get_gpu_device()
    test_data = DataSource(device, CurveKind.Polynomials)
    net = Net(n_feature=1, n_hidden=20, n_output=1).to(device)

    #optimizer是训练的工具
    optimizer = torch.optim.SGD(net.parameters(), lr=0.01)  #传入net的所有参数，学习率
    loss_func = torch.nn.MSELoss()  #预测值和真实值的误差计算公式（均方差）

    for t in range(2000):
        optimizer.zero_grad()  #清空上一步的残余更新参数值
        prediction: torch.tensor = net(test_data.x)  #喂给net训练数据x，输出预测值
        loss: torch.tensor = loss_func(prediction, test_data.y)  #计算两者的误差
        loss.backward()  #误差反向传播，计算参数更新值
        optimizer.step()  #将参数更新值施加到net的parameters上

    #plot and show learning process
    plt.cla()
    plt.scatter(test_data.x.cpu().data.numpy(), test_data.y.cpu().data.numpy())
    plt.scatter(test_data.x.cpu().data.numpy(), prediction.cpu().data.numpy())
    plt.text(0.5,
             0,
             'Loss=%.4f' % loss.cpu().data.numpy(),
             fontdict={
                 'size': 20,
                 'color': 'red'
             })
    plt.show()


if __name__ == "__main__":
    gpu_utils.print_gpu_availability()
    main()