import torch
import numpy as np
import pandas as pd
from model import CardiovascularRNN
from sklearn.preprocessing import StandardScaler, LabelEncoder

class RiskPredictor:
    def __init__(self, model_path='cardiovascular_rnn_model.pth', csv_path='cardiovascular_risk_dataset.csv'):
        self.df = pd.read_csv(csv_path)
        
        # Determine feature names and order
        self.feature_names = self.df.drop(['Patient_ID', 'risk_category'], axis=1).columns.tolist()
        
        self.le_risk = LabelEncoder()
        self.le_risk.fit(self.df['risk_category'])
        
        self.le_smoking = LabelEncoder()
        self.le_smoking.fit(self.df['smoking_status'])
        
        self.le_family = LabelEncoder()
        self.le_family.fit(self.df['family_history_heart_disease'])
        
        self.scaler = StandardScaler()
        df_proc = self.df.drop(['Patient_ID', 'risk_category'], axis=1)
        df_proc['smoking_status'] = self.le_smoking.transform(df_proc['smoking_status'])
        df_proc['family_history_heart_disease'] = self.le_family.transform(df_proc['family_history_heart_disease'])
        self.scaler.fit(df_proc.values)
        
        self.input_size = 1
        self.hidden_size = 64
        self.num_layers = 2
        self.num_classes = len(self.le_risk.classes_)
        
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = CardiovascularRNN(self.input_size, self.hidden_size, self.num_layers, self.num_classes).to(self.device)
        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
        self.model.eval()
        print(f"Model initialized with features: {self.feature_names}")

    def predict_single(self, data_dict):
        # Convert dictionary to DataFrame with explicit column order
        input_df = pd.DataFrame([data_dict])[self.feature_names]
        
        # Transform categorical
        input_df['smoking_status'] = self.le_smoking.transform(input_df['smoking_status'])
        input_df['family_history_heart_disease'] = self.le_family.transform(input_df['family_history_heart_disease'])
        
        # Scale
        input_scaled = self.scaler.transform(input_df.values)
        input_tensor = torch.FloatTensor(input_scaled).reshape(1, -1, 1).to(self.device)
        
        # Predict
        with torch.no_grad():
            output = self.model(input_tensor)
            probs = torch.softmax(output, dim=1)
            _, predicted = torch.max(output, 1)
            predicted_label = self.le_risk.inverse_transform([predicted.item()])[0]
            
        print(f"Prediction: {predicted_label} | Probabilities: {probs.cpu().numpy()}")
        return predicted_label

def predict():
    predictor = RiskPredictor()
    print("Model loaded successfully.")

    # Take a sample from the dataset for prediction
    df = pd.read_csv('cardiovascular_risk_dataset.csv')
    sample_row = df.drop(['Patient_ID', 'risk_category'], axis=1).iloc[0]
    sample_dict = sample_row.to_dict()
    true_label = df['risk_category'].iloc[0]
    
    predicted_label = predictor.predict_single(sample_dict)
    
    print(f"\nSample Data: {sample_dict}")
    print(f"True Risk Category: {true_label}")
    print(f"Predicted Risk Category: {predicted_label}")

if __name__ == "__main__":
    predict()