import joblib
import pandas as pd
import numpy as np

def load_model(path):
    return joblib.load(path)

def predict_single(model, sample_dict):
    """
    sample_dict: dict mapping column name -> value for the input features (do NOT include target)
    """
    X = pd.DataFrame([sample_dict])
    proba = None
    try:
        proba = model.predict_proba(X)
    except Exception:
        pass
    pred = model.predict(X)
    return {"prediction": pred[0], "probability": proba.tolist() if proba is not None else None}

if __name__ == "__main__":
    # Example usage with Kidney Disease model
    model_path_kidney = "/content/models/kidney_model.pkl" # Path to the saved kidney model
    model_kidney = load_model(model_path_kidney)

    # Sample data for Kidney Disease prediction
    # Include all columns from the merged dataset, fill with relevant values or None/np.nan
    sample_kidney = {
        'Age': None,
        'Gender': None,
        'Total_Bilirubin': None,
        'Direct_Bilirubin': None,
        'Alkaline_Phosphotase': None,
        'Alamine_Aminotransferase': None,
        'Aspartate_Aminotransferase': None,
        'Total_Protiens': None,
        'Albumin': None,
        'Albumin_and_Globulin_Ratio': None,
        'Dataset': None, # Target column for liver, set to None
        'id': None, # ID column, set to None
        'age': 48.0, # Example kidney feature
        'bp': 80.0, # Example kidney feature
        'sg': 1.020, # Example kidney feature
        'al': 1.0, # Example kidney feature
        'su': 0.0, # Example kidney feature
        'rbc': 'NaN', # Example kidney feature (categorical)
        'pc': 'normal', # Example kidney feature (categorical)
        'pcc': 'notpresent', # Example kidney feature (categorical)
        'ba': 'notpresent', # Example kidney feature (categorical)
        'bgr': 121.0, # Example kidney feature
        'bu': 36.0, # Example kidney feature
        'sc': 1.2, # Example kidney feature
        'sod': np.nan, # Example kidney feature (missing value)
        'pot': np.nan, # Example kidney feature (missing value)
        'hemo': 15.4, # Example kidney feature
        'pcv': 44.0, # Example kidney feature
        'wc': 7800.0, # Example kidney feature
        'rc': 5.2, # Example kidney feature
        'htn': 'yes', # Example kidney feature (categorical)
        'dm': 'yes', # Example kidney feature (categorical)
        'cad': 'no', # Example kidney feature (categorical)
        'appet': 'good', # Example kidney feature (categorical)
        'pe': 'no', # Example kidney feature (categorical)
        'ane': 'no', # Example kidney feature (categorical)
        'classification': None, # Target column for kidney, set to None
        'name': None, # Name column for parkinsons, set to None
        "MDVP:Fo(Hz)": None, # Parkinsons feature, set to None
        "MDVP:Fhi(Hz)": None, # Parkinsons feature, set to None
        "MDVP:Flo(Hz)": None, # Parkinsons feature, set to None
        "MDVP:Jitter(%)": None,
        "MDVP:Jitter(Abs)": None,
        "MDVP:RAP": None,
        "MDVP:PPQ": None,
        "Jitter:DDP": None,
        "MDVP:Shimmer": None,
        "MDVP:Shimmer(dB)": None,
        "Shimmer:APQ3": None,
        "Shimmer:APQ5": None,
        "MDVP:APQ": None,
        "Shimmer:DDA": None,
        "NHR": None,
        "HNR": None,
        "status": None, # Target column for parkinsons, set to None
        "RPDE": None,
        "DFA": None,
        "spread1": None,
        "spread2": None,
        "D2": None,
        "PPE": None
    }
    res_kidney = predict_single(model_kidney, sample_kidney)
    print("Kidney Disease Prediction:", res_kidney)

    # Example usage with Liver Patient model
    model_path_liver = "/content/models/liver_model.pkl"
    model_liver = load_model(model_path_liver)
    # Sample data for Liver Patient prediction
    # Include all columns from the merged dataset, fill with relevant values or None/np.nan
    sample_liver = {
        'Age': 65.0, # Example liver feature
        'Gender': 'Female', # Example liver feature (categorical)
        'Total_Bilirubin': 0.7, # Example liver feature
        'Direct_Bilirubin': 0.1, # Example liver feature
        'Alkaline_Phosphotase': 187.0, # Example liver feature
        'Alamine_Aminotransferase': 16.0, # Example liver feature
        'Aspartate_Aminotransferase': 18.0, # Example liver feature
        'Total_Protiens': 6.8, # Example liver feature
        'Albumin': 3.3, # Example liver feature
        'Albumin_and_Globulin_Ratio': 0.90, # Example liver feature
        'Dataset': None, # Target column for liver, set to None
        'id': None, # ID column, set to None
        'age': None, # Kidney feature, set to None
        'bp': None, # Kidney feature, set to None
        'sg': None, # Kidney feature, set to None
        'al': None, # Kidney feature, set to None
        'su': None, # Kidney feature, set to None
        'rbc': None, # Kidney feature, set to None
        'pc': None, # Kidney feature, set to None
        'pcc': None, # Kidney feature, set to None
        'ba': None, # Kidney feature, set to None
        'bgr': None, # Kidney feature, set to None
        'bu': None, # Kidney feature, set to None
        'sc': None, # Kidney feature, set to None
        'sod': None, # Kidney feature, set to None
        'pot': None, # Kidney feature, set to None
        'hemo': None, # Kidney feature, set to None
        'pcv': None, # Kidney feature, set to None
        'wc': None, # Kidney feature, set to None
        'rc': None, # Kidney feature, set to None
        'htn': None, # Kidney feature, set to None
        'dm': None, # Kidney feature, set to None
        'cad': None, # Kidney feature, set to None
        'appet': None, # Kidney feature, set to None
        'pe': None, # Kidney feature, set to None
        'ane': None, # Kidney feature, set to None
        'classification': None, # Target column for kidney, set to None
        'name': None, # Name column for parkinsons, set to None
        "MDVP:Fo(Hz)": None, # Parkinsons feature, set to None
        "MDVP:Fhi(Hz)": None, # Parkinsons feature, set to None
        "MDVP:Flo(Hz)": None, # Parkinsons feature, set to None
        "MDVP:Jitter(%)": None,
        "MDVP:Jitter(Abs)": None,
        "MDVP:RAP": None,
        "MDVP:PPQ": None,
        "Jitter:DDP": None,
        "MDVP:Shimmer": None,
        "MDVP:Shimmer(dB)": None,
        "Shimmer:APQ3": None,
        "Shimmer:APQ5": None,
        "MDVP:APQ": None,
        "Shimmer:DDA": None,
        "NHR": None,
        "HNR": None,
        "status": None, # Target column for parkinsons, set to None
        "RPDE": None,
        "DFA": None,
        "spread1": None,
        "spread2": None,
        "D2": None,
        "PPE": None
    }
    res_liver = predict_single(model_liver, sample_liver)
    print("Liver Patient Prediction:", res_liver)