import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, r2_score, make_scorer, mean_absolute_error
from sklearn.neural_network import MLPRegressor
import pickle

# データの読み込みと準備
X = pd.read_excel('defect_data_all.xlsx')
Y = pd.read_excel('TFT_data_all.xlsx')

# データをトレーニングセットとテストセットに分割
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=42)

# 標準化
sc_X = StandardScaler()
X_train_scaled = sc_X.fit_transform(X_train)
X_test_scaled = sc_X.transform(X_test)
sc_Y = StandardScaler()
Y_train_scaled = sc_Y.fit_transform(Y_train)
Y_test_scaled = sc_Y.transform(Y_test)

# PCAによる次元削減と標準化
pca = PCA(n_components=4)
Y_train_pca = pca.fit_transform(Y_train_scaled)
Y_test_pca = pca.transform(Y_test_scaled)
sc_Y_pca = StandardScaler()
Y_train_pca_scaled = sc_Y_pca.fit_transform(Y_train_pca)
Y_test_pca_scaled = sc_Y_pca.transform(Y_test_pca)

# PyTorchのテンソルに変換
X_train_tensor = torch.FloatTensor(X_train_scaled)
Y_train_pca_tensor = torch.FloatTensor(Y_train_pca_scaled)
X_test_tensor = torch.FloatTensor(X_test_scaled)
Y_test_tensor = torch.FloatTensor(Y_test_pca_scaled)

# データローダーの作成
#train_data = TensorDataset(X_train_tensor, Y_train_pca_tensor)
#train_loader = DataLoader(train_data, batch_size=64, shuffle=True)

# 多層パーセプトロンの定義
class MLPRegressor_torch(nn.Module):
    def __init__(self, input_size, output_size):
        super(MLPRegressor_torch, self).__init__()
        self.fc1 = nn.Linear(input_size, 10)
        self.fc2 = nn.Linear(10, 10)
        self.fc3 = nn.Linear(10, 10)
        self.fc4 = nn.Linear(10, 10)
        self.fc5 = nn.Linear(10, output_size)
        
    def forward(self, x):
        x = torch.tanh(self.fc1(x))
        x = torch.tanh(self.fc2(x))
        x = torch.tanh(self.fc3(x))
        x = torch.tanh(self.fc4(x))
        x = self.fc5(x)
        return x


model_inverse = MLPRegressor_torch(input_size=Y_train_pca.shape[1], output_size=X_train.shape[1])
optimizer = optim.LBFGS(model_inverse.parameters())

# 損失関数の定義
criterion = nn.MSELoss()

#順モデルの読込み
model_forward = pickle.load(open('mlp_Y-model.sav', 'rb'))

#Loss_XとLoss_Yの重み
R_X = 1
R_Y = 1

#L2正則化パラメータ
alpha = 0.01

# トレーニング
num_epochs = 100
for epoch in range(num_epochs):
    model_inverse.train()
    optimizer.zero_grad()
    X_cal_train = model_inverse(Y_train_pca)
    Y_cal_train_array = model_forward.predict(X_cal_train.detach().numpy())
    Y_cal_train = torch.FloatTensor(Y_cal_train_array)
    
    loss_X = criterion(X_cal_train, X_train)
    loss_Y = criterion(Y_cal_train, Y_train)

    l2 = torch.tensor(0., requires_grad=True)
    for w in model_inverse.parameters():
        l2 = l2 + torch.norm(w)**2

    loss = R_X * loss_X + R_Y * loss_Y + alpha * l2

    loss.backward()
    optimizer.step()

    # エポックごとに損失を表示
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}')

# テストデータでの評価
model_inverse.eval()
with torch.no_grad():
    X_cal_test = model_inverse(Y_test_tensor)
    Y_cal_test_array = model_forward.predict(X_cal_test.numpy())

    mae = mean_absolute_error(Y_test, Y_cal_test_array)
    r2 = r2_score(Y_test, Y_cal_test_array)

print(f'MAE: {mae}')
print(f'R2: {r2}')