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

# データの読み込み
df = pd.read_csv('distribution/datasets/dataset.csv')
for col in df.columns:
    if col == "SMILES":
        continue
    df[col] = df[col].fillna(0)

# λmax以外の列を特徴量として選択
X = df.drop(['SMILES', 'λmax'], axis=1).values.astype(float)

# λmaxを目的変数として選択
y = df['λmax'].values.astype(float)

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

# 標準化
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# PCAによる次元削減
pca = PCA(n_components=15)
X_train_pca = pca.fit_transform(X_train_scaled)
X_test_pca = pca.transform(X_test_scaled)

# PyTorchのテンソルに変換
X_train_tensor = torch.FloatTensor(X_train_pca)
y_train_tensor = torch.FloatTensor(y_train).view(-1, 1)
X_test_tensor = torch.FloatTensor(X_test_pca)
y_test_tensor = torch.FloatTensor(y_test).view(-1, 1)

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

# 多層パーセプトロンの定義
class MLPRegressor(nn.Module):
    def __init__(self, input_size):
        super(MLPRegressor, self).__init__()
        self.fc1 = nn.Linear(input_size, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 32)
        self.fc4 = nn.Linear(32, 16)
        self.fc5 = nn.Linear(16, 1)
        
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        x = self.fc5(x)
        return x


model = MLPRegressor(input_size=X_train_pca.shape[1])
optimizer = optim.Adam(model.parameters())

scorer = make_scorer(mean_squared_error, greater_is_better=False)

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

# トレーニング
num_epochs = 100
for epoch in range(num_epochs):
    model.train()
    for inputs, targets in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

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

# テストデータでの評価
model.eval()
with torch.no_grad():
    y_pred = model(X_test_tensor)
    mse = criterion(y_pred, y_test_tensor)
    r2 = r2_score(y_test, y_pred.numpy())

print(f'Mean Squared Error: {mse.item()}')
print(f'R-squared: {r2}')