# -*- coding: utf-8 -*-
"""
@author: Mohammad Mohammadiaria
"""
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Transformer–LSTM Model Training Script for Multi-Stressor Time-Series Prediction of Malignancy.

This Python script implements the hybrid Transformer–LSTM model described in the manuscript, designed for
 time-series classification of cell malignancy based on mechanistic simulation outputs. 
 The model integrates transformer-based self-attention for capturing long-range dependencies 
 with LSTM sequence modeling for temporal dynamics. The script supports .npz input datasets 
 containing multi-stressor profiles (ROS, pH, temperature, ion channel conductances, etc.) 
 and outputs predicted malignancy classes.
"""

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.metrics import accuracy_score, roc_auc_score, confusion_matrix
import pandas as pd
npz_path = r"C:\MDA_timeseries.npz"
# ----------------------------
# Model Definition
# ----------------------------

class TransformerLSTM(nn.Module):
    def __init__(self, input_dim, d_model=64, nhead=4, num_layers=2, lstm_hidden=64, num_classes=2):
        super(TransformerLSTM, self).__init__()
        self.input_fc = nn.Linear(input_dim, d_model)
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.lstm = nn.LSTM(d_model, lstm_hidden, batch_first=True)
        self.fc_out = nn.Linear(lstm_hidden, num_classes)

    def forward(self, x):
        x = self.input_fc(x)  # (batch, seq_len, d_model)
        x = self.transformer_encoder(x)
        lstm_out, _ = self.lstm(x)
        out = self.fc_out(lstm_out[:, -1, :])  # take last time step
        return out

# ----------------------------
# Load Data
# ----------------------------

def load_npz(path):
    data = np.load(path)
    X = data['X']  # shape: (samples, timesteps, features)
    y = data['y']  # shape: (samples,)
    return X, y

# ----------------------------
# Train Function
# ----------------------------

def train_model(X, y, num_epochs=50, batch_size=16, lr=1e-3):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    n_samples, seq_len, n_features = X.shape

    model = TransformerLSTM(input_dim=n_features).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    dataset = TensorDataset(torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.long))
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    model.train()
    for epoch in range(num_epochs):
        epoch_loss = 0
        for X_batch, y_batch in loader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            optimizer.zero_grad()
            outputs = model(X_batch)
            loss = criterion(outputs, y_batch)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()

        print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss/len(loader):.4f}")

    return model

# ----------------------------
# Evaluation
# ----------------------------

def evaluate_model(model, X, y):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.eval()
    with torch.no_grad():
        outputs = model(torch.tensor(X, dtype=torch.float32).to(device))
        preds = torch.argmax(outputs, dim=1).cpu().numpy()
        acc = accuracy_score(y, preds)
        auc = roc_auc_score(y, preds)
        cm = confusion_matrix(y, preds)
    return acc, auc, cm, preds

# ----------------------------
# Main
# ----------------------------

if __name__ == "__main__":
    # Path to NPZ time-series dataset (change this to your file)
    npz_path = "MDA_timeseries.npz"
    X, y = load_npz(npz_path)

    model = train_model(X, y, num_epochs=50, batch_size=16, lr=1e-3)

    acc, auc, cm, preds = evaluate_model(model, X, y)
    print("Accuracy:", acc)
    print("AUC:", auc)
    print("Confusion Matrix:\n", cm)

    # Save predictions to Excel
    df_preds = pd.DataFrame({
        "y_true": y,
        "y_pred": preds
    })
    df_preds.to_excel("transformer_lstm_results.xlsx", index=False)
    print("Results saved to transformer_lstm_results.xlsx")