import torch
import numpy as np

class LSTMPredictor(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, forecast_horizon, n_layers=1, dropout=0.0):
        super(LSTMPredictor, self).__init__()
        self.lstm = torch.nn.LSTM(input_dim, hidden_dim, n_layers, dropout=dropout, batch_first=True)
        self.fc = torch.nn.Linear(hidden_dim, output_dim)
        self.forecast_horizon = forecast_horizon

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        forecast = self.fc(lstm_out[:, -self.forecast_horizon:, :])  # Only take last forecast_horizon time steps
        return forecast

# Load the model
def load_model(model_path):
    model_state = torch.load(model_path)
    model = LSTMPredictor(
        input_dim=model_state['model_architecture']['input_dim'],
        hidden_dim=model_state['model_architecture']['hidden_dim'],
        output_dim=model_state['model_architecture']['output_dim'],
        forecast_horizon=model_state['model_architecture']['forecast_horizon'],
        n_layers=model_state['model_architecture']['n_layers'],
        dropout=model_state['model_architecture']['dropout']
    )
    model.load_state_dict(model_state['model_state_dict'])
    
    # Optionally load any other necessary parameters (like scaling if saved separately)
    return model

# Inference function
def predict(features):
    # Convert input features to tensor
    input_tensor = torch.FloatTensor(features)
    
    # Get model prediction
    model.eval()
    with torch.no_grad():
        predictions = model(input_tensor).numpy()
    
    # Return the predictions
    return predictions.tolist()

# Main function to load the model and make predictions
if __name__ == "__main__":
    
    model = load_model('lstm_crisis_severity_predictor_20241116_092126.pt')  # Path to your model file

    # Example input data (features)
    test_features = np.array([[1.23, 4.56, 7.89, 10.11]])  # Example test features, replace with actual data

    # Get predictions
    predictions = predict(test_features)
    
    # Output predictions
    print("Predictions:", predictions)