# FNN

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as dsets
import pandas as pd
from normal_use import *

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class FNN_Net(nn.Module):

    def __init__(self):
        super(FNN_Net, self).__init__()
        self.features = 0
        self.linear_relu1 = nn.Linear(self.features, 128)
        self.linear_relu2 = nn.Linear(128, 256)
        self.linear_relu3 = nn.Linear(256, 256)
        self.linear_relu4 = nn.Linear(256, 256)
        self.linear5 = nn.Linear(256, 1)

    def forward(self, x):
        y_pred = self.linear_relu1(x)
        y_pred = nn.functional.relu(y_pred)

        y_pred = self.linear_relu2(y_pred)
        y_pred = nn.functional.relu(y_pred)

        y_pred = self.linear_relu3(y_pred)
        y_pred = nn.functional.relu(y_pred)

        y_pred = self.linear_relu4(y_pred)
        y_pred = nn.functional.relu(y_pred)

        y_pred = self.linear5(y_pred)
        return y_pred

    def fit(self, X, y_all):
        """
        Description
        -----------
        fit

        Parameters
        ----------
        X - train data X
        y_all - train data ys
        
        """

        y = y_all
        # 训练数据集特征
        train_features = torch.tensor(X.values, dtype=torch.float)
        train_features = train_features.to(device)
        # 训练数据集目标
        train_labels = torch.tensor(y.values, dtype=torch.float).view(-1, 1)
        train_labels = train_labels.to(device)

        self.features=train_features.shape[1]
        self.linear_relu1 = nn.Linear(self.features, 128)
        self = self.to(device)
        self.train()

        criterion = nn.MSELoss(reduction='mean')
        criterion = criterion.to(device)
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-4)

        losses = []
        epoch = 5000
        epoch_real = 0
        for t in range(epoch):
            y_pred = self(train_features)
            y_pred.to(device)

            loss = criterion(y_pred, train_labels)
            losses.append(loss.item())

            if torch.isnan(loss):
                break

            optimizer.zero_grad()

            loss.backward()

            optimizer.step()
            epoch_real = epoch_real + 1

        print(f"epoch:{epoch_real} loss:{losses[-1]}")


    def predict(self, X):
        self.to(device)
        test_features = torch.tensor(X.values, dtype=torch.float)
        test_features = test_features.to(device)
        pred_labels = self(test_features)
        pred_labels = pred_labels.cpu().data.numpy()
        return pred_labels


# # DNN equal next NN
# import itertools
# from sklearn.model_selection import train_test_split
# import pandas as pd
# import numpy as np
# from keras import layers, models

# # 模型定义

# class DNN_Net():
#     def build_model(self):
#         network = models.Sequential()
#         network.add(layers.Dense(64, activation='relu', input_shape=(13, )))
#         network.add(layers.Dense(64, activation='relu'))
#         network.add(layers.Dense(1))  # 最后输出预测值，恒等函数
#         #损失函数用mes(均方误差), 监控指标为mae(平均绝对误差, 返回误差绝对值)
#         network.compile(optimizer='rmsprop', loss='mse', metrics=['mae'])
#         return network

#     def __init__(self):
#         self.network = self.build_model()

    
#     def fit(self, X, ys):
#         self.network.summary()
#         self.network.fit(X, ys, epochs=1, batch_size=1)

#     def predict(self, X):
#         res = self.network.predict(X)
#         print(res)
#         return res


# NN

import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
from sklearn import metrics

class NN_Net(nn.Module):
    def __init__(self):
        super(NN_Net, self).__init__()
        self.linear_relu1 = nn.Linear(13, 64)
        self.relu6 = nn.ReLU6()
        # self.relu6 = nn.ReLU()
        self.linear_relu2 = nn.Linear(64, 64)
        self.leaky = nn.LeakyReLU()
        # self.leaky = nn.ReLU()
        self.linear3 = nn.Linear(64, 1)

    def forward(self, x):
        y_pred = self.linear_relu1(x)
        y_pred = self.relu6(y_pred)

        y_pred = self.linear_relu2(y_pred)
        y_pred = self.leaky(y_pred)

        y_pred = self.linear3(y_pred)
        return y_pred

    def fit(self, X, y_all):
        """
        Description
        -----------
        fit

        Parameters
        ----------
        X - train data X
        y_all - train data ys
        
        """
        # 训练数据集特征
        train_features = torch.tensor(X.values, dtype=torch.float)
        train_features = train_features.to(device)
        # 训练数据集目标
        train_labels = torch.tensor(y_all.values, dtype=torch.float).view(-1, 1)
        train_labels = train_labels.to(device)

        self = self.to(device)
        # self.train()

        opti = torch.optim.SGD(self.parameters(), lr=0.02)
        loss_func = nn.MSELoss()
        train_dataloader = DataLoader(TensorDataset(train_features, train_labels), batch_size=320)
        for t in range(1000):
            for batch, (x, y) in enumerate(train_dataloader):
                pred = self(x)
                loss = loss_func(pred, y)
                opti.zero_grad()
                torch.sqrt(loss).backward()
                opti.step()
            print(t)


    def predict(self, X):
        self.to(device)
        test_features = torch.tensor(X.values, dtype=torch.float)
        test_features = test_features.to(device)
        pred_labels = self(test_features)
        pred_labels = pred_labels.cpu().data.numpy()
        return pred_labels