import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error
import joblib

# 1. Load and preprocess data
dfs = []
for i in range(1, 12):
    try:
        df = pd.read_csv(f'SAMPLE{i}.csv')
        df.columns = df.columns.str.strip()
        dfs.append(df)
    except FileNotFoundError:
        print(f"SAMPLE{i}.csv not found")

df = pd.concat(dfs, ignore_index=True)

columns_to_drop = ['cid', 'cmpdname', 'cmpdsynonym', 'inchi', 'smiles', 'inchikey', 'iupacname',
                   'meshheadings', 'annothits', 'annotation', 'mf', 'aids', 'cidcdate', 'sidsrcname', 'depcatg']
X = df.drop(columns=columns_to_drop, axis=1, errors='ignore')
target_columns = ['mw', 'polararea']
target_columns = [col for col in target_columns if col in df.columns]
y = df[target_columns].apply(pd.to_numeric, errors='coerce')

numeric_cols = X.select_dtypes(include=['number']).columns
for col in numeric_cols:
    X[col] = pd.to_numeric(X[col], errors='coerce')
    X[col] = X[col].fillna(X[col].mean())
X = X.select_dtypes(include=['number'])

X = X[y.notna().all(axis=1)]
y = y[y.notna().all(axis=1)]

X_train, X_test, y_train, y_test = train_test_split(X.values, y.values, test_size=0.2, random_state=42)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

joblib.dump(scaler, "scaler.h5")

# 2. PyTorch Dataset
class DrugDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.float32)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

train_dataset = DrugDataset(X_train_scaled, y_train)
test_dataset = DrugDataset(X_test_scaled, y_test)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 3. Model
class MultiOutputRegressorNN(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, output_dim)
        )

    def forward(self, x):
        return self.network(x)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MultiOutputRegressorNN(input_dim=X.shape[1], output_dim=len(target_columns)).to(device)

# 4. Train
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(100):
    model.train()
    total_loss = 0
    for X_batch, y_batch in train_loader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)
        optimizer.zero_grad()
        output = model(X_batch)
        loss = criterion(output, y_batch)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}, Loss: {total_loss / len(train_loader):.4f}")

# 5. Evaluate
model.eval()
all_preds, all_targets = [], []
with torch.no_grad():
    for X_batch, y_batch in test_loader:
        X_batch = X_batch.to(device)
        preds = model(X_batch).cpu().numpy()
        all_preds.append(preds)
        all_targets.append(y_batch.numpy())

y_pred = np.vstack(all_preds)
y_true = np.vstack(all_targets)

for i, col in enumerate(target_columns):
    rmse = mean_squared_error(y_true[:, i], y_pred[:, i], squared=False)
    print(f"RMSE for {col}: {rmse:.2f}")

# 6. Save Model
torch.save(model.state_dict(), "drug_prediction_model.pth")
print("Model saved as drug_prediction_model.pth")

# 7. Predict New Sample
new_data = pd.DataFrame({
    'polararea': [17.1],
    'complexity': [0.0],
    'xlogp': [-0.32],
    'heavycnt': [2],
    'hbonddonor': [1],
    'hbondacc': [1],
    'rotbonds': [0],
    'exactmass': [46.0689],
    'monoisotopicmass': [46.0689],
    'charge': [0],
    'covalentunitcnt': [1]
})

for col in X.columns:
    if col not in new_data.columns:
        new_data[col] = 0
new_data = new_data[X.columns]
new_data_scaled = scaler.transform(new_data)
new_tensor = torch.tensor(new_data_scaled, dtype=torch.float32).to(device)

model.eval()
with torch.no_grad():
    prediction = model(new_tensor).cpu().numpy()[0]

for i, col in enumerate(target_columns):
    print(f"Predicted {col}: {prediction[i]:.2f}")