import pandas as pd
from transformers import BertTokenizer, BertForSequenceClassification
import torch

# Step 1: Load the dataset (adjust the dataset path if necessary)
df = pd.read_json("hf://datasets/theprint/mindfulness-alpaca/alpaca_data_export.json")

# Check the dataset structure to understand how the data is formatted
print(df.head())  # Ensure it contains columns like 'health_issue' or related info

# Step 2: Load the pre-trained BERT model and tokenizer
model = BertForSequenceClassification.from_pretrained("bert-base-uncased")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

# Function to predict the health issue based on user input (using BERT)
def predict_health_issue(question):
    inputs = tokenizer(question, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=-1).item()
    return predicted_class

# Step 3: Define a mapping of health issues to suggested cures or workouts
health_cures = {
    "fatigue": "Try regular cardio workouts and improving your sleep habits.",
    "joint pain": "Consider low-impact exercises like swimming or yoga.",
    "stress": "Meditate regularly and try mindfulness exercises.",
    "headaches": "Stay hydrated and manage your screen time.",
    "weight gain": "Focus on a balanced diet and regular exercise like walking or jogging.",
    "insomnia": "Develop a consistent bedtime routine and try relaxation techniques."
}

# Step 4: Ask the user a series of health-related questions and process responses
questions = [
    "Do you feel constant fatigue?",
    "Have you been experiencing any joint pain?",
    "Do you have difficulty sleeping?",
    "Are you feeling stressed frequently?",
    "Have you been gaining weight unexpectedly?",
    "Do you experience frequent headaches?"
]

responses = []

print("Please answer the following questions to help identify your health issues:")

# Ask the questions and record responses
for question in questions:
    response = input(question + " (Yes/No/Other): ")
    responses.append(response)

# Step 5: Predict health issues and suggest cures
print("\nHealth Issue Analysis:")
for response in responses:
    predicted_issue = predict_health_issue(response)
    health_issue = df['health_issue'].iloc[predicted_issue]  # Map prediction to a health issue
    print(f"Predicted Health Issue: {health_issue}")

    if health_issue in health_cures:
        print(f"Suggested cure for {health_issue}: {health_cures[health_issue]}")
    else:
        print(f"No suggestion available for {health_issue}. Please consult a healthcare professional.")