"""
Interactive Grocery List Recommender (ML-Powered)
Uses best-performing Random Forest model trained on real Pakistan food prices.
"""

import json
from pathlib import Path
import joblib

import numpy as np
import pandas as pd

# Import recommendation layer for ML-based ranking
from recommendation_layer import (
    compute_recommendation_scores,
    recommend_products,
    get_recommendation_summary
)

# Paths
BASE_DIR = Path(__file__).resolve().parent
PAKISTAN_PRICES_PATH = BASE_DIR / "Pakistan_Food_Prices_2025.csv"
GROCERY_CATALOG_PATH = BASE_DIR / "Grocery_data (1).csv"
MODELS_DIR = BASE_DIR / "models"

# Global model and encoders (loaded once)
DATASET = None
MODEL = None


def _patch_column_transformer(model):
    """Ensure ColumnTransformer has expected attrs across sklearn versions."""
    try:
        pre = getattr(model, "named_steps", {}).get("pre")
        if pre is not None and not hasattr(pre, "_name_to_fitted_passthrough"):
            pre._name_to_fitted_passthrough = {}
    except Exception:
        pass


def _build_catalog_from_pakistan_prices(df: pd.DataFrame) -> pd.DataFrame:
    """Aggregate Pakistan price dataset into an item catalog with median prices."""
    # Clean columns
    df = df.copy()
    df['Item'] = df['Item'].astype(str).str.strip()
    df['Category'] = df['Category'].astype(str).str.strip()
    df['Price_per_Kg'] = pd.to_numeric(df['Price_per_Kg'], errors='coerce')
    df = df.dropna(subset=['Item', 'Category', 'Price_per_Kg'])

    # Aggregate median price per Item + Category
    agg = (
        df.groupby(['Item', 'Category'], as_index=False)['Price_per_Kg']
          .median()
          .rename(columns={'Price_per_Kg': 'price'})
    )
    return agg


def load_dataset():
    """Load real datasets and prepare an item catalog for recommendations."""
    global DATASET

    if not PAKISTAN_PRICES_PATH.exists():
        raise FileNotFoundError(f"Dataset not found: {PAKISTAN_PRICES_PATH}")

    pak_df = pd.read_csv(PAKISTAN_PRICES_PATH)
    catalog_df = _build_catalog_from_pakistan_prices(pak_df)

    # Optional: merge with grocery catalog when available (future enhancement)
    DATASET = catalog_df.drop_duplicates(subset=['Item']).reset_index(drop=True)

    print(f"[OK] Dataset loaded: {len(DATASET)} items from Pakistan price data")


def load_models():
    """Load unified diet compatibility model trained on real data."""
    global MODEL
    path = MODELS_DIR / "diet_unified_model.joblib"
    if not path.exists():
        raise FileNotFoundError(f"Model not found: {path}. Please run: python train_model.py")
    MODEL = joblib.load(path)
    _patch_column_transformer(MODEL)
    print("[OK] Unified diet model loaded")


def predict_diet_compatibility(target_diet: str):
    """Compute recommendation scores using ML diet probability blended with price."""
    df = DATASET.copy()
    # Pipeline expects columns: Item, Category, price, DietType
    diet = 'Normal' if target_diet == 'All' else target_diet
    df = df.copy()
    df['DietType'] = diet
    # Lasso returns continuous predictions; clip to [0, 1] for probability interpretation
    model_pred = MODEL.predict(df[["Item", "Category", "price", "DietType"]])
    model_proba = np.clip(model_pred, 0, 1)

    scored_df = compute_recommendation_scores(df[["Item", "Category", "price"]], diet_type=diet, ml_diet_proba=model_proba)
    return scored_df


def generate_grocery_list(budget: float, family_size: int, diet_type: str) -> dict:
    """Generate grocery list with ML-based ranking and separation of purchased vs recommended
    
    Returns:
        Dictionary containing:
        - purchased_items: Items selected within budget (what user should buy)
        - recommended_items: High-scoring items not purchased (suggestions for consideration)
        - budget_summary: Cost details
        
    Why this separation:
        - Purchased items = actionable shopping list
        - Recommended items = ML-powered suggestions to help users discover alternatives
        - Keeps output focused while providing intelligent recommendations on demand
    """
    
    # Get all products with ML-based recommendation scores
    products_scored = predict_diet_compatibility(diet_type)
    
    # Filter by recommendation score threshold (0.2 = reasonable confidence)
    products_filtered = products_scored[products_scored['recommendation_score'] > 0.2].copy()
    
    if len(products_filtered) == 0:
        # Fallback: use top 20% by recommendation score
        threshold = products_scored['recommendation_score'].quantile(0.80)
        products_filtered = products_scored[products_scored['recommendation_score'] >= threshold].copy()
    
    # Use recommendation layer to rank and select items within budget
    purchased_items, total_cost = recommend_products(products_filtered, budget, family_size)
    
    # Get purchased product names for filtering
    purchased_names = {item['product'] for item in purchased_items}

    # Determine the product name column in dataframe
    name_col = 'Item' if 'Item' in products_filtered.columns else 'product'

    # Get recommended items (high-scoring but NOT purchased)
    # These are suggestions the user may consider
    recommended_items = products_filtered[
        ~products_filtered[name_col].isin(purchased_names)
    ].copy()
    
    # Sort recommended items by recommendation score (best first)
    recommended_items = recommended_items.sort_values(
        by='recommendation_score',
        ascending=False
    ).head(15)  # Limit to top 15 recommendations
    
    # Format recommended items for display
    recommended_list = []
    for _, row in recommended_items.iterrows():
        recommended_list.append({
            'product': row.get('Item', row.get('product')),
            'category': row.get('Category', row.get('category')),
            'price': float(row['price']),
            'recommendation_score': float(row['recommendation_score']),
            'diet_suitability': float(row.get('diet_score', np.nan)),
            'price_affordability': float(row['price_score'])
        })
    
    # Calculate budget summary
    remaining = budget - total_cost
    
    result = {
        'budget_pkr': float(budget),
        'family_size': int(family_size),
        'diet_type': diet_type,
        'total_cost': round(total_cost, 2),
        'remaining': round(remaining, 2),
        'purchased_items': purchased_items,
        'recommended_items': recommended_list
    }
    
    return result


def get_recommended_items(result: dict) -> list:
    """Extract recommended items from result dictionary
    
    These are ML-ranked products that scored well but weren't purchased
    due to budget constraints. Useful for:
    - Discovering alternatives
    - Planning future purchases
    - Understanding what the ML model considers suitable
    """
    return result.get('recommended_items', [])


def display_results(result: dict):
    """Display purchased grocery list (actionable shopping list)
    
    This is what the user should actually buy - clean and focused.
    Recommendations are shown separately only on request.
    """
    print("\n" + "="*60)
    print("  🛒 FINAL GROCERY LIST (Items to Buy)")
    print("="*60)
    
    print(f"\n📋 Budget Details:")
    print(f"   Total Budget: PKR {result['budget_pkr']:,.2f}")
    print(f"   Family Size: {result['family_size']} members")
    print(f"   Diet Type: {result['diet_type']}")
    print(f"   Total Cost: PKR {result['total_cost']:,.2f}")
    print(f"   Remaining: PKR {result['remaining']:,.2f}")
    
    # Get recommendation summary statistics
    purchased = result['purchased_items']
    if purchased:
        rec_scores = [item['recommendation_score'] for item in purchased]
        avg_score = np.mean(rec_scores)
        
        print(f"\n📊 Selection Quality:")
        print(f"   Items Selected: {len(purchased)}")
        print(f"   Avg ML Score: {avg_score:.2%}")
    
    # Display items by category
    print(f"\n🛍️  Shopping List:")
    print("-" * 60)
    
    # Group by category
    from collections import defaultdict
    by_category = defaultdict(list)
    for item in purchased:
        by_category[item['category']].append(item)
    
    # Sort categories for consistent display
    for category in sorted(by_category.keys()):
        items = by_category[category]
        cat_total = sum(item['total_cost'] for item in items)
        print(f"\n{category.upper()} (PKR {cat_total:,.2f})")
        
        for item in items:
            print(f"  • {item['product']:<35} PKR {item['total_cost']:>8,.2f}")
    
    print("\n" + "="*60)


def show_recommended_items(result: dict):
    """Display ML-ranked recommendations (items user may consider)
    
    These are high-scoring products NOT purchased due to budget constraints.
    Shown only when user explicitly requests them.
    
    Why show these:
    - Help users discover alternatives
    - Understand what ML considers suitable for their diet
    - Plan future purchases or substitutions
    """
    recommended = result.get('recommended_items', [])
    
    if not recommended:
        print("\n💡 No additional recommendations available at this time.")
        return
    
    print("\n" + "="*60)
    print("  💡 RECOMMENDED ITEMS (You May Consider)")
    print("="*60)
    print("\nThese are high-quality alternatives suggested by ML:")
    print(f"Total: {len(recommended)} items\n")
    
    # Display with detailed scoring
    print(f"{'Product':<35} {'Price':>10} {'ML Score':>10} {'Details':>20}")
    print("-" * 80)
    
    for item in recommended:
        product = item['product'][:33]  # Truncate if too long
        price = item['price']
        score = item['recommendation_score']
        diet_suit = item['diet_suitability']
        price_afford = item['price_affordability']
        
        # Create details string
        details = f"D:{diet_suit:.0%} P:{price_afford:.0%}"
        
        print(f"{product:<35} PKR {price:>7,.2f} {score:>9.1%}  {details:>20}")
    
    print("\n" + "="*60)
    print("Legend: ML Score = Overall recommendation strength")
    print("        D = Diet Suitability (60% weight)")
    print("        P = Price Affordability (40% weight)")
    print("="*60)


def display_json(result: dict):
    """Display results as JSON"""
    print("\n" + "="*80)
    print("JSON OUTPUT")
    print("="*80)
    
    # Create simplified output for JSON
    json_output = {
        'budget_pkr': result['budget_pkr'],
        'family_size': result['family_size'],
        'diet_type': result['diet_type'],
        'total_cost': result['total_cost'],
        'remaining': result['remaining'],
        'purchased_items': result['purchased_items'],
        'recommended_items': result['recommended_items']
    }
    
    print(json.dumps(json_output, indent=2))
    print("="*80)


def prompt_budget(default: float = 5000) -> float:
    """Prompt user for budget in PKR"""
    while True:
        try:
            print()
            response = input(f">>> Enter your budget in PKR [default: {default}]: ").strip()
            if not response:
                print(f"[Selected] Budget: {default} PKR")
                return float(default)
            budget = float(response)
            print(f"[Selected] Budget: {budget} PKR")
            return budget
        except ValueError:
            print("[ERROR] Invalid input. Please enter a valid number.")



def prompt_family_size(default: int = 3) -> int:
    """Prompt user for family size"""
    while True:
        try:
            print()
            response = input(f">>> Enter family size (number of people) [default: {default}]: ").strip()
            if not response:
                print(f"[Selected] Family Size: {default} people")
                return int(default)
            size = int(response)
            if size < 1:
                print("[ERROR] Family size must be at least 1.")
                continue
            print(f"[Selected] Family Size: {size} people")
            return size
        except ValueError:
            print("[ERROR] Invalid input. Please enter a valid number.")


def prompt_diet_type(default: str = "Normal") -> str:
    """Prompt user for diet type"""
    valid_diets = ["Normal", "Diabetic", "Keto", "All", "Vegetarian"]
    
    print()
    print("Available diet types:")
    for i, diet in enumerate(valid_diets, 1):
        print(f"  {i}. {diet}")
    
    while True:
        print()
        response = input(f">>> Select diet type (enter name or number) [default: {default}]: ").strip()
        
        if not response:
            print(f"[Selected] Diet Type: {default}")
            return default
        
        # Try numeric selection
        try:
            choice = int(response)
            if 1 <= choice <= len(valid_diets):
                selected = valid_diets[choice - 1]
                print(f"[Selected] Diet Type: {selected}")
                return selected
            else:
                print(f"[ERROR] Please enter a number between 1 and {len(valid_diets)}.")
                continue
        except ValueError:
            pass
        
        # Try text selection (case-insensitive)
        selected = response.capitalize()
        if selected in valid_diets:
            print(f"[Selected] Diet Type: {selected}")
            return selected
        
        print(f"[ERROR] Invalid diet. Please choose: {', '.join(valid_diets)}")


def main():
    print("\n" + "="*80)
    print("[GROCERY LIST RECOMMENDER]")
    print("Module: AI-powered recommendations using Random Forest (Real Dataset)")
    print("="*80)
    print("\nThis system recommends items based on:")
    print("  * Your budget")
    print("  * Family size")
    print("  * Diet type")
    
    # Load model and data
    print("\n[Loading real datasets...]")
    try:
        load_dataset()
        load_models()
    except FileNotFoundError as e:
        print(f"\nERROR: {e}")
        return
    
    # Get user input
    print("\n" + "="*80)
    print("STEP 1: ENTER YOUR PREFERENCES")
    print("="*80)
    
    budget = prompt_budget()
    family_size = prompt_family_size()
    diet_type = prompt_diet_type()
    
    # Generate grocery list with ML-based recommendations
    print("\n" + "="*80)
    print("STEP 2: GENERATING GROCERY LIST")
    print("="*80)
    print("\n[Using ML model to rank products and select optimal items...]")
    result = generate_grocery_list(
        budget=budget,
        family_size=family_size,
        diet_type=diet_type
    )
    
    # Display purchased items (always shown)
    print("\n" + "="*80)
    print("STEP 3: YOUR GROCERY LIST")
    print("="*80)
    display_results(result)
    
    # Ask if user wants to see recommendations (shown only on request)
    print("\n" + "="*80)
    print("STEP 4: ADDITIONAL OPTIONS")
    print("="*80)
    response = input("\nDo you want to see recommended items? (yes/no) [no]: ").strip().lower()
    if response in ['yes', 'y']:
        show_recommended_items(result)
    
    # Ask if user wants JSON output
    response = input("\nDo you want to export JSON? (yes/no) [no]: ").strip().lower()
    if response in ['yes', 'y']:
        display_json(result)
    
    # Ask if user wants to generate another list
    response = input("\nGenerate another list? (yes/no) [no]: ").strip().lower()
    if response in ['yes', 'y']:
        main()
    else:
        print("\nThank you for using ML Grocery Recommender!")
        print("="*80 + "\n")


if __name__ == "__main__":
    main()