import pandas as pd
import json
import os
from sklearn.model_selection import train_test_split

# ✅ Path to Yelp dataset (Update if needed)
data_path = "../yelp_academic_dataset_review.json"

# ✅ List to store processed rows
rows = []

# ✅ Read JSON file (JSON Lines format)
print("📥 Loading and processing dataset...")
with open(data_path, "r", encoding="utf-8") as f:
    for line in f:
        review = json.loads(line)
        rows.append({
            "text": review["text"], 
            "stars": int(review["stars"]) - 1,  # Convert 1-5 stars to 0-4 for training
            "useful": int(round(min(max(review["useful"], 1), 5))),  # Normalize usefulness (1-5)
            "response": review.get("response", ""),  # Keep provided responses, empty if missing
        })

# ✅ Convert to Pandas DataFrame
df = pd.DataFrame(rows)

# ✅ Reduce dataset size for quick testing (Use full dataset for actual training)
df = df.sample(100000, random_state=42)  # Sample 100K reviews for testing

# ✅ Train/Test Split (80/20 split)
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

# ✅ Ensure 'data' directory exists
os.makedirs("data", exist_ok=True)

# ✅ Save processed data as CSV
train_df.to_csv("../../data/train.csv", index=False)
test_df.to_csv("../../data/test.csv", index=False)

print("✅ Data successfully saved to 'data/train.csv' and 'data/test.csv'!")