clean_and_enrich.py

import csv
import json
import os
import re
import time
import litellm
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()


# Litellm model configurations using the Nvidia OpenAI-compatible endpoint
MODEL = "openai/" + os.environ.get("ANTHROPIC_DEFAULT_SONNET_MODEL", "moonshotai/kimi-k2.6")
API_KEY = os.environ.get("ANTHROPIC_API_KEY")
BASE_URL = os.environ.get("ANTHROPIC_BASE_URL")

# Timings Parsing Logic
DAYS_MAP = {
    'mon': 0, 'tue': 1, 'wed': 2, 'thu': 3, 'fri': 4, 'sat': 5, 'sun': 6,
    'monday': 0, 'tuesday': 1, 'wednesday': 2, 'thursday': 3, 'friday': 4, 'saturday': 5, 'sunday': 6
}

def parse_day_range(day_str):
    day_str = day_str.strip().lower()
    if ',' in day_str:
        days = []
        for part in day_str.split(','):
            days.extend(parse_day_range(part))
        return list(set(days))
        
    if '-' in day_str:
        parts = day_str.split('-')
        if len(parts) == 2:
            start_name = parts[0].strip()
            end_name = parts[1].strip()
            if start_name in DAYS_MAP and end_name in DAYS_MAP:
                start_idx = DAYS_MAP[start_name]
                end_idx = DAYS_MAP[end_name]
                if start_idx <= end_idx:
                    return list(range(start_idx, end_idx + 1))
                else: # wrap around, e.g. Fri-Tue -> 4,5,6,0,1,2
                    return list(range(start_idx, 7)) + list(range(0, end_idx + 1))
    
    if day_str in DAYS_MAP:
        return [DAYS_MAP[day_str]]
    return []

def time_to_minutes(t_str, is_closing=False):
    t_str = t_str.strip().lower()
    if not t_str:
        return None
        
    if t_str == '12noon' or t_str == 'noon':
        return 12 * 60
    if t_str == '12midnight' or t_str == 'midnight':
        return 24 * 60 if is_closing else 0
        
    match = re.match(r'(\d+)(?::?(\d{2}))?\s*(am|pm|noon|midnight)?', t_str)
    if not match:
        return None
        
    hour_str, min_str, mer = match.groups()
    hour = int(hour_str)
    
    if min_str:
        minute = int(min_str)
    else:
        if len(hour_str) >= 3:
            minute = int(hour_str[-2:])
            hour = int(hour_str[:-2])
        else:
            minute = 0
            
    if mer:
        if mer == 'pm' and hour < 12:
            hour += 12
        elif mer == 'am' and hour == 12:
            hour = 0
        elif mer == 'noon':
            hour = 12
            minute = 0
        elif mer == 'midnight':
            hour = 24 if is_closing else 0
            minute = 0
            
    if is_closing and hour < 12 and (mer == 'am' or not mer):
        hour += 24
        
    return hour * 60 + minute

def parse_timings(timings_str):
    schedule = {i: [] for i in range(7)}
    t_str = timings_str.strip()
    if not t_str or t_str.upper() in ('NA', 'OPENING'):
        return {i: [[0, 1440]] for i in range(7)}
        
    if t_str == '24':
        return {i: [[0, 1440]] for i in range(7)}
        
    if t_str.endswith('...'):
        t_str = t_str[:-3].strip()
        
    pattern = r'([^()]+)(?:\(([^()]+)\))?'
    matches = re.findall(pattern, t_str)
    
    specified_days = set()
    temp_rules = []
    
    for time_part, day_part in matches:
        time_part = time_part.strip().strip(',')
        if not time_part:
            continue
            
        days = parse_day_range(day_part) if day_part else []
        is_closed = 'closed' in time_part.lower()
        
        intervals = []
        if not is_closed:
            sub_parts = time_part.split(',')
            for sub in sub_parts:
                sub = sub.strip()
                if ' to ' in sub:
                    times = sub.split(' to ')
                    if len(times) == 2:
                        start_min = time_to_minutes(times[0], is_closing=False)
                        end_min = time_to_minutes(times[1], is_closing=True)
                        if start_min is not None and end_min is not None:
                            intervals.append([start_min, end_min])
        
        temp_rules.append((days, intervals, is_closed))
        if days:
            specified_days.update(days)
            
    for days, intervals, is_closed in temp_rules:
        target_days = days if days else [d for d in range(7) if d not in specified_days]
        if not days and not specified_days:
            target_days = list(range(7))
            
        for d in target_days:
            if is_closed:
                schedule[d] = []
            else:
                schedule[d].extend(intervals)
                
    for d in range(7):
        if not schedule[d] and d not in specified_days:
            schedule[d] = [[0, 1440]]
            
    return schedule

# Rule-based fallback generator
def generate_fallback_qualitative(row):
    vibe = "Casual"
    good_for = ["friends", "family"]
    has_parking = "street"
    noise_level = "moderate"
    must_try_dish = ""
    
    cuisines_str = row.get("cuisines", "").lower()
    r_type = row.get("type", "").lower()
    cost = int(row.get("cost_for_two", 0)) if row.get("cost_for_two") else 0
    
    # 1. Type specific vibe and must try
    if r_type == 'cafe':
        vibe = "Cafe"
        good_for = ["friends", "couples", "solo"]
        must_try_dish = "Cold Coffee"
    elif 'street food' in cuisines_str or 'chaat' in cuisines_str:
        vibe = "Street Food"
        good_for = ["friends", "solo"]
        noise_level = "lively"
        must_try_dish = "Pani Puri"
    elif 'desserts' in cuisines_str or 'bakery' in cuisines_str or 'ice cream' in cuisines_str:
        vibe = "Dessert Parlour"
        good_for = ["family", "friends", "couples"]
        must_try_dish = "Special Ice Cream"
    elif 'fast food' in cuisines_str or 'pizza' in cuisines_str or 'burger' in cuisines_str:
        vibe = "Fast Food"
        good_for = ["friends", "family"]
        must_try_dish = "Cheese Burger"
    # Cuisines check
    elif 'gujarati' in cuisines_str or 'kathiyawadi' in cuisines_str:
        vibe = "Casual"
        good_for = ["family", "large groups"]
        noise_level = "lively"
        must_try_dish = "Gujarati Thali"
        has_parking = "yes"
    elif 'punjabi' in cuisines_str or 'north indian' in cuisines_str:
        must_try_dish = "Paneer Tikka Masala"
        good_for = ["family", "large groups", "friends"]
    elif 'south indian' in cuisines_str:
        must_try_dish = "Masala Dosa"
        good_for = ["family", "solo"]
    else:
        must_try_dish = "Special Platter"

    # 2. Cost specific overrides
    if cost >= 1000:
        vibe = "Fine Dining"
        good_for = ["couples", "family", "business"]
        noise_level = "quiet"
        has_parking = "yes"
    elif cost >= 600:
        has_parking = "yes"
        
    return {
        "vibe": vibe,
        "good_for": json.dumps(good_for),
        "has_parking": has_parking,
        "noise_level": noise_level,
        "must_try_dish": must_try_dish
    }

# LLM batch enricher
def enrich_batch_with_llm(restaurants_batch):
    prompt_input = []
    for r in restaurants_batch:
        prompt_input.append({
            "id": r["detail_url"] or r["name"],
            "name": r["name"],
            "locality": r["locality"],
            "cuisines": r["cuisines"],
            "type": r["type"],
            "cost_for_two": r["cost_for_two"]
        })
        
    system_prompt = (
        "You are an expert local food guide specializing in Ahmedabad's food scene.\n"
        "Your task is to enrich the given list of restaurants with qualitative characteristics based on their name, locality, cuisines, type, and cost.\n"
        "Output must be a JSON array containing objects with these exact keys:\n"
        "- id: the id provided in the input\n"
        "- vibe: choose one of [Casual, Fine Dining, Street Food, Cafe, Rooftop, Dessert Parlour, Fast Food, Dhaba]\n"
        "- good_for: list containing one or more of [family, couples, friends, solo, business, large groups]\n"
        "- has_parking: choose one of [yes, no, street]\n"
        "- noise_level: choose one of [quiet, moderate, lively]\n"
        "- must_try_dish: free text string representing a famous or standard signature dish there (keep it under 4 words, e.g. 'Cheese Butter Masala', 'Cold Coffee', 'Gujarati Thali').\n\n"
        "Output ONLY the JSON list. Do not write any markdown tags (like ```json), introduction, or footer."
    )
    
    user_prompt = f"Restaurants to enrich:\n{json.dumps(prompt_input, indent=2)}"
    
    try:
        response = litellm.completion(
            model=MODEL,
            api_key=API_KEY,
            base_url=BASE_URL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            temperature=0.0
        )
        content = response.choices[0].message.content.strip()
        
        # Clean any accidental markdown code fences
        if content.startswith("```"):
            lines = content.split("\n")
            if lines[0].startswith("```"):
                lines = lines[1:]
            if lines[-1].strip() == "```":
                lines = lines[:-1]
            content = "\n".join(lines).strip()
            
        data = json.loads(content)
        # Create lookup dictionary
        enriched_map = {}
        for item in data:
            # good_for needs to be serialized as json string
            good_for_list = item.get("good_for", ["family", "friends"])
            if isinstance(good_for_list, str):
                good_for_list = [good_for_list]
            enriched_map[item["id"]] = {
                "vibe": item.get("vibe", "Casual"),
                "good_for": json.dumps(good_for_list),
                "has_parking": item.get("has_parking", "street"),
                "noise_level": item.get("noise_level", "moderate"),
                "must_try_dish": item.get("must_try_dish", "")
            }
        return enriched_map
    except Exception as e:
        print(f"Error calling LLM for batch: {e}")
        return {}

def main():
    input_file = "ahmedabad_every_food_place.csv"
    output_file = "ahmedabad_restaurants_enriched.csv"
    
    print(f"Starting data cleaning pipeline for {input_file}...")
    
    # 1. Read input rows
    rows = []
    with open(input_file, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for r in reader:
            rows.append(r)
            
    print(f"Total raw rows: {len(rows)}")
    
    # 2. Clean base fields
    cleaned_rows = []
    for r in rows:
        # Clean cost_for_two: e.g. "₹250 for two" -> 250
        cost_str = r.get("cost_for_two", "").replace("₹", "")
        cost_match = re.search(r'\d+', cost_str)
        cost_val = int(cost_match.group()) if cost_match else None
        
        # Clean user_avg_rating: e.g. "-" or "NEW" or "Opening" -> None
        rating_str = r.get("user_avg_rating", "").strip()
        rating_val = None
        if rating_str not in ("-", "NEW", "Opening", "N/A", ""):
            try:
                rating_val = float(rating_str)
            except ValueError:
                pass
                
        # Clean user_rating_count: e.g. "-" or "NEW" -> 0
        votes_str = r.get("user_rating_count", "").strip()
        votes_val = 0
        if votes_str not in ("-", "NEW", "Opening", "N/A", ""):
            try:
                votes_val = int(votes_str)
            except ValueError:
                pass
                
        # Parse timings to structured JSON string
        timings_raw = r.get("timings", "")
        parsed_timings_dict = parse_timings(timings_raw)
        timings_json = json.dumps(parsed_timings_dict)
        
        cleaned_row = {
            "name": r.get("name", "").strip(),
            "type": r.get("type", "restaurant").strip(),
            "place_area": r.get("place_area", "").strip(),
            "street_address": r.get("street_address", "").strip(),
            "locality": r.get("locality", "").strip(),
            "city": "ahmedabad",
            "user_avg_rating": rating_val,
            "user_rating_count": votes_val,
            "cost_for_two": cost_val,
            "cuisines": r.get("cuisines", "").strip(),
            "timings": timings_raw,
            "timings_json": timings_json,
            "detail_url": r.get("detail_url", "").strip()
        }
        cleaned_rows.append(cleaned_row)
        
    print("Base fields cleaned successfully.")
    
    # 3. Check for checkpoint (already enriched records)
    enriched_cache = {}
    if os.path.exists(output_file):
        print(f"Checkpoint file found at {output_file}. Loading cache to skip completed records...")
        try:
            with open(output_file, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                for r in reader:
                    key = r.get("detail_url") or r.get("name")
                    if r.get("vibe"): # Has qualitative fields
                        enriched_cache[key] = {
                            "vibe": r.get("vibe"),
                            "good_for": r.get("good_for"),
                            "has_parking": r.get("has_parking"),
                            "noise_level": r.get("noise_level"),
                            "must_try_dish": r.get("must_try_dish")
                        }
            print(f"Loaded {len(enriched_cache)} cached records.")
        except Exception as e:
            print(f"Error loading checkpoint: {e}")
            
    # 4. Identify which rows should get LLM enrichment vs Rule fallback
    # Filter condition for top-rated restaurants: rating >= 3.8 and at least 10 votes
    llm_enrich_queue = []
    fallback_queue = []
    
    for r in cleaned_rows:
        key = r["detail_url"] or r["name"]
        if key in enriched_cache:
            continue
            
        rating = r["user_avg_rating"]
        votes = r["user_rating_count"]
        if rating is not None and rating >= 3.8 and votes >= 10:
            llm_enrich_queue.append(r)
        else:
            fallback_queue.append(r)
            
    print(f"Queue size: LLM enrichment = {len(llm_enrich_queue)}, Rule-based fallback = {len(fallback_queue)}")
    
    # 5. Process LLM queue in batches of 50
    enriched_results = enriched_cache
    batch_size = 50
    
    if llm_enrich_queue:
        print(f"Starting LLM enrichment for {len(llm_enrich_queue)} restaurants in batches of {batch_size}...")
        for idx in range(0, len(llm_enrich_queue), batch_size):
            batch = llm_enrich_queue[idx:idx + batch_size]
            print(f"Enriching batch {idx//batch_size + 1} of {(len(llm_enrich_queue) - 1)//batch_size + 1}...")
            
            # Retrieve enriched values
            enriched_batch_map = enrich_batch_with_llm(batch)
            
            # Merge into overall results
            for r in batch:
                key = r["detail_url"] or r["name"]
                if key in enriched_batch_map:
                    enriched_results[key] = enriched_batch_map[key]
                else:
                    # Fail-safe rule fallback if LLM batch call missed or failed for this record
                    enriched_results[key] = generate_fallback_qualitative(r)
            
            # Save checkpoint intermediate results to prevent work loss
            # Write current snapshot of all rows to file
            write_dataset(output_file, cleaned_rows, enriched_results)
            time.sleep(0.5) # rate limit friendly delay
            
    # 6. Apply fallback rules for the rest
    if fallback_queue:
        print(f"Applying rule-based defaults for {len(fallback_queue)} restaurants...")
        for r in fallback_queue:
            key = r["detail_url"] or r["name"]
            enriched_results[key] = generate_fallback_qualitative(r)
            
    # 7. Final write of dataset
    print(f"Saving final dataset to {output_file}...")
    write_dataset(output_file, cleaned_rows, enriched_results)
    print("Done! Data cleaning and enrichment complete.")

def write_dataset(filename, cleaned_rows, enriched_results):
    fieldnames = [
        "name", "type", "place_area", "street_address", "locality", "city", 
        "user_avg_rating", "user_rating_count", "cost_for_two", "cuisines", 
        "timings", "timings_json", "detail_url", "vibe", "good_for", 
        "has_parking", "noise_level", "must_try_dish"
    ]
    
    with open(filename, 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        
        for r in cleaned_rows:
            key = r["detail_url"] or r["name"]
            q_info = enriched_results.get(key, {
                "vibe": "Casual",
                "good_for": json.dumps(["family", "friends"]),
                "has_parking": "street",
                "noise_level": "moderate",
                "must_try_dish": ""
            })
            
            row_out = {**r, **q_info}
            writer.writerow(row_out)

if __name__ == "__main__":
    main()