import torch.nn as nn
import matplotlib.pyplot as plt
import torch
from torch import optim
import numpy as np
torch.manual_seed(1)    # 保證程序隨機生成數一樣
x1 = torch.rand(200) * 2
x2 = torch.rand(200) * 2
data = zip(x1,x2)
pos = []                # 定義類型 1
neg = []                # 定義類型 2
def classification(data):
    for i in data:
        if(i[1] - i[0] < 0):
            pos.append(i)
        else:
            neg.append(i)
classification(data)
pos_x = [i[0] for i in pos]
pos_y = [i[1] for i in pos]
neg_x = [i[0] for i in neg]
neg_y = [i[1] for i in neg]
plt.scatter(pos_x,pos_y,c='r')
plt.scatter(neg_x,neg_y,c='b')
plt.show()
x_data = [[i[0],i[1]] for i in pos]
x_data.extend([[i[0],i[1]] for i in neg])
x_data = torch.Tensor(x_data)     # 輸入數據 feature
y_data = [1 for i in range(len(pos))]
y_data.extend([0 for i in range(len(neg))])
y_data = torch.Tensor(y_data).view(-1,1)    # 對應的標簽
class LogisticRegressionModel(nn.Module):   # 定義網絡
    def __init__(self):
        super(LogisticRegressionModel,self).__init__()
        self.linear = nn.Linear(2,1)
        self.sigmoid = nn.Sigmoid()
    def forward(self,x):
        x = self.linear(x)
        x = self.sigmoid(x)
        return x
model = LogisticRegressionModel()
criterion = nn.BCELoss()    
optimizer = optim.SGD(model.parameters(),lr =0.01)
for epoch in range(10000):
    y_pred = model(x_data)
    loss = criterion(y_pred,y_data)     # 計算損失值
    if epoch % 1000 == 0:
        print(epoch,loss.item())     # 打印損失值
    optimizer.zero_grad()         # 梯度清零
    loss.backward()                 # 反向傳播
    optimizer.step()                # 梯度更新
w = model.linear.weight[0]      # 取出訓練完成的結果
w0 = w[0]
w1 = w[1]
b = model.linear.bias.item()
with torch.no_grad():      # 繪制決策邊界，這里不需要計算梯度
    x= torch.arange(0,3).view(-1,1)
    y = (- w0 * x - b) / w1
    plt.plot(x.numpy(),y.numpy())
plt.scatter(pos_x,pos_y,c='r')          
plt.scatter(neg_x,neg_y,c='b')
plt.xlim(0,2)
plt.ylim(0,2)
plt.show()