#!/usr/bin/env python3
"""
Interactive Crop Yield Predictor

This script provides an interactive interface to predict crop yields using trained models.
Users can input parameters and get predictions from all three models (Random Forest, XGBoost, PyTorch).
"""

import pandas as pd
import numpy as np
import joblib
import warnings
import os
import json
import sys
from datetime import datetime

# Important: import the training module so that the pickled DataPreprocessor
# class can be resolved during joblib.load()
import crop_yield_ml_pipeline  # noqa: F401

warnings.filterwarnings('ignore')

class DataPreprocessor:
    """Data preprocessing and feature engineering class."""
    
    def __init__(self):
        self.label_encoders = {}
        self.scaler = None
        self.imputer = None
        self.feature_names = None
        
    def prepare_features(self, df):
        """Prepare features for machine learning."""
        # Create a copy to avoid modifying original data
        data = df.copy()
        
        # Feature engineering
        data['Area_Production_Ratio'] = data['Area'] / (data['Production'] + 1e-6)
        data['Yield_Area_Interaction'] = data.get('Yield', 0) * data['Area']
        data['Production_Per_Area'] = data['Production'] / (data['Area'] + 1e-6)
        
        # Create season dummies
        season_dummies = pd.get_dummies(data['Season'], prefix='Season')
        
        # Add missing season columns with zeros if they don't exist
        expected_seasons = ['Season_Autumn', 'Season_Kharif', 'Season_Rabi', 
                          'Season_Summer', 'Season_Total', 'Season_Whole Year', 'Season_Winter']
        for season in expected_seasons:
            if season not in season_dummies.columns:
                season_dummies[season] = 0
        
        data = pd.concat([data, season_dummies[expected_seasons]], axis=1)
        
        # Handle categorical variables
        categorical_cols = ['State', 'District', 'Crop']
        
        for col in categorical_cols:
            if col in data.columns and col in self.label_encoders:
                # Handle unseen categories
                unique_values = set(data[col].astype(str))
                known_values = set(self.label_encoders[col].classes_)
                new_values = unique_values - known_values
                
                if new_values:
                    # For new categories, assign them the most common category's code
                    mode_value = self.label_encoders[col].classes_[0]
                    data[col] = data[col].astype(str).replace(list(new_values), mode_value)
                
                data[f'{col}_encoded'] = self.label_encoders[col].transform(data[col].astype(str))
            elif col in data.columns:
                # If encoder doesn't exist, use simple integer encoding
                unique_vals = data[col].astype(str).unique()
                data[f'{col}_encoded'] = pd.Categorical(data[col].astype(str)).codes
        
        # Select features for modeling
        feature_cols = ['Crop_Year', 'Area', 'Production', 'Annual_Rainfall', 
                       'Fertilizer', 'Pesticide', 'State_encoded', 'Crop_encoded',
                       'Area_Production_Ratio', 'Yield_Area_Interaction', 
                       'Production_Per_Area'] + expected_seasons
        
        # Add District_encoded if available
        if 'District_encoded' in data.columns:
            feature_cols.append('District_encoded')
        
        # Select only available columns
        available_cols = [col for col in feature_cols if col in data.columns]
        X = data[available_cols].copy()
        
        return X, data
    
    def transform(self, X):
        """Transform new data using fitted preprocessors."""
        if self.imputer is None or self.scaler is None:
            raise ValueError("Preprocessor not fitted. Please load a trained preprocessor.")
        
        # Handle missing values
        X_imputed = pd.DataFrame(
            self.imputer.transform(X), 
            columns=X.columns, 
            index=X.index
        )
        
        # Scale features
        X_scaled = pd.DataFrame(
            self.scaler.transform(X_imputed), 
            columns=X.columns, 
            index=X.index
        )
        
        return X_scaled

# PyTorch and XGBoost models removed - using only Random Forest for simplicity

class CropYieldPredictor:
    """Main prediction class that loads Random Forest model and makes predictions."""
    
    def __init__(self, models_dir='trained_models', quiet=False):
        self.models_dir = models_dir
        self.model = None
        self.preprocessor = None
        self.quiet = quiet
        
        if not quiet:
            print(f"🚀 Initializing Random Forest Crop Yield Predictor...")
        
        self.load_models()
        
    def load_models(self):
        """Load Random Forest model and preprocessor."""
        if not self.quiet:
            print("📥 Loading trained model...")
        
        try:
            # Load preprocessor
            preprocessor_path = os.path.join(self.models_dir, 'preprocessor.pkl')
            if os.path.exists(preprocessor_path):
                self.preprocessor = joblib.load(preprocessor_path)
                if not self.quiet:
                    print("  ✅ Preprocessor loaded")
            else:
                raise FileNotFoundError("Preprocessor not found. Please train models first.")
            
            # Load Random Forest
            rf_path = os.path.join(self.models_dir, 'random_forest_model.pkl')
            if os.path.exists(rf_path):
                self.model = joblib.load(rf_path)
                if not self.quiet:
                    print("  ✅ Random Forest model loaded")
            else:
                raise FileNotFoundError("Random Forest model not found. Please train models first.")
                
        except Exception as e:
            if not self.quiet:
                print(f"❌ Error loading models: {e}")
            raise
    
    def predict_yield(self, input_data):
        """Make yield prediction using Random Forest model."""
        
        try:
            # Convert input to DataFrame
            if isinstance(input_data, dict):
                df = pd.DataFrame([input_data])
            else:
                df = input_data.copy()
            
            # Prepare features
            X, processed_data = self.preprocessor.prepare_features(df)
            
            # Transform data
            X_processed = self.preprocessor.transform(X)
            
            # Make prediction with Random Forest
            try:
                prediction = self.model.predict(X_processed)[0]
                prediction = max(0, prediction)  # Ensure non-negative yield
                return prediction, processed_data
                
            except Exception as e:
                return f"Error: {str(e)}", None
            
        except Exception as e:
            return f"Error: {str(e)}", None
    
    def get_crop_options(self):
        """Get available crop options from the preprocessor."""
        if 'Crop' in self.preprocessor.label_encoders:
            return list(self.preprocessor.label_encoders['Crop'].classes_)
        return []
    
    def get_state_options(self):
        """Get available state options from the preprocessor."""
        if 'State' in self.preprocessor.label_encoders:
            return list(self.preprocessor.label_encoders['State'].classes_)
        return []
    
    def get_season_options(self):
        """Get available season options."""
        return ['Kharif', 'Rabi', 'Summer', 'Whole Year', 'Autumn', 'Winter', 'Total']

def interactive_prediction():
    """Interactive command-line interface for yield prediction."""
    
    print("=" * 70)
    print("🌾 CROP YIELD PREDICTION SYSTEM 🌾")
    print("=" * 70)
    
    # Initialize predictor
    try:
        predictor = CropYieldPredictor()
        print(f"\n✅ System ready! Random Forest model loaded.")
        print("\n📋 Available options:")
        print(f"   States: {len(predictor.get_state_options())} available")
        print(f"   Crops: {len(predictor.get_crop_options())} available")
        print(f"   Seasons: {len(predictor.get_season_options())} available")
    except Exception as e:
        print(f"❌ Failed to initialize predictor: {e}")
        print("Please ensure you have trained models by running: python crop_yield_ml_pipeline.py")
        return
    
    while True:
        print("\n" + "=" * 70)
        print("🔍 ENTER PREDICTION PARAMETERS")
        print("=" * 70)
        
        try:
            # Get input parameters
            print("📅 Basic Information:")
            crop_year = int(input("  Crop Year (e.g., 2024): "))
            
            print("\n🌾 Crop and Location:")
            state = input("  State (e.g., 'Punjab', 'Uttar Pradesh'): ").strip()
            district = input("  District (optional, press Enter to skip): ").strip() or "Unknown"
            crop = input("  Crop (e.g., 'Rice', 'Wheat', 'Maize'): ").strip()
            season = input("  Season (Kharif/Rabi/Summer/Whole Year): ").strip()
            
            print("\n📊 Agricultural Data:")
            area = float(input("  Area (in hectares): "))
            production = float(input("  Production (in tons): "))
            
            print("\n🌧️ Environmental & Input Data (optional - press Enter to use defaults):")
            rainfall_input = input("  Annual Rainfall (mm, default=1000): ").strip()
            annual_rainfall = float(rainfall_input) if rainfall_input else 1000.0
            
            fertilizer_input = input("  Fertilizer usage (kg, default=50): ").strip()
            fertilizer = float(fertilizer_input) if fertilizer_input else 50.0
            
            pesticide_input = input("  Pesticide usage (kg, default=5): ").strip()
            pesticide = float(pesticide_input) if pesticide_input else 5.0
            
            # Create input data
            input_data = {
                'Crop_Year': crop_year,
                'State': state,
                'District': district,
                'Crop': crop,
                'Season': season,
                'Area': area,
                'Production': production,
                'Annual_Rainfall': annual_rainfall,
                'Fertilizer': fertilizer,
                'Pesticide': pesticide
            }
            
            print("\n🔄 Processing prediction...")
            
            # Make prediction
            prediction, processed_data = predictor.predict_yield(input_data)
            
            # Display results
            print("\n" + "=" * 70)
            print("🎯 YIELD PREDICTION RESULTS")
            print("=" * 70)
            
            if isinstance(prediction, str) and "Error" in prediction:
                print(f"❌ {prediction}")
            else:
                print(f"📍 Input Summary:")
                print(f"   📅 Year: {crop_year}")
                print(f"   🌾 Crop: {crop} ({season} season)")
                print(f"   📍 Location: {district}, {state}")
                print(f"   📏 Area: {area} hectares")
                print(f"   📦 Production: {production} tons")
                print(f"   🌧️ Rainfall: {annual_rainfall} mm")
                print(f"   🌱 Fertilizer: {fertilizer} kg")
                print(f"   🧪 Pesticide: {pesticide} kg")
                
                print(f"\n🎯 Predicted Yield:")
                print(f"   Random Forest: {prediction:8.2f} kg/hectare")
                
                # Calculate total expected production
                total_production = (prediction * area) / 1000  # Convert to tons
                print(f"\n📦 Total Expected Production: {total_production:.2f} tons")
                
                # Provide interpretation
                print(f"\n💡 Interpretation:")
                if prediction > 3000:
                    print("   🟢 Excellent yield expected!")
                elif prediction > 2000:
                    print("   🟡 Good yield expected.")
                elif prediction > 1000:
                    print("   🟠 Moderate yield expected.")
                else:
                    print("   🔴 Low yield expected. Consider optimization.")
            
        except KeyboardInterrupt:
            print("\n\n👋 Goodbye!")
            break
        except ValueError as e:
            print(f"❌ Invalid input: {e}")
        except Exception as e:
            print(f"❌ Error during prediction: {e}")
        
        # Ask if user wants to continue
        print("\n" + "-" * 70)
        continue_choice = input("🔄 Make another prediction? (y/n): ").strip().lower()
        if continue_choice not in ['y', 'yes']:
            print("\n👋 Thank you for using the Crop Yield Prediction System!")
            break

def format_json_output(prediction, area, assessment_text):
    """Format prediction results as JSON output."""
    total_production = (prediction * area) / 1000.0  # Convert to tons
    
    # Extract assessment without emoji
    assessment_map = {
        "🟢 Excellent yield expected!": "Excellent yield expected",
        "🟡 Good yield expected.": "Good yield expected",
        "🟠 Moderate yield expected.": "Moderate yield expected",
        "🔴 Low yield expected.": "Low yield expected"
    }
    
    clean_assessment = assessment_map.get(assessment_text, assessment_text)
    
    result = {
        "model": "Random Forest",
        "predicted_yield": f"{round(prediction, 2)} kg/hectare",
        "total_expected_production": f"{round(total_production, 2)} tons",
        "assessment": clean_assessment
    }
    
    return result

def validate_json_input(data):
    """Validate and normalize JSON input data."""
    required_fields = ['year', 'state', 'crop', 'season', 'area', 'production']
    optional_fields = {'rainfall': 1000.0, 'fertilizer': 50.0, 'pesticide': 5.0}
    
    # Check required fields
    for field in required_fields:
        if field not in data:
            raise ValueError(f"Missing required field: {field}")
        if data[field] is None or data[field] == "":
            raise ValueError(f"Field '{field}' cannot be empty")
    
    # Add optional fields with defaults
    for field, default_value in optional_fields.items():
        if field not in data or data[field] is None:
            data[field] = default_value
    
    # Convert to internal format
    input_data = {
        'Crop_Year': int(data['year']),
        'State': str(data['state']),
        'District': "Unknown",  # Default district
        'Crop': str(data['crop']),
        'Season': str(data['season']),
        'Area': float(data['area']),
        'Production': float(data['production']),
        'Annual_Rainfall': float(data['rainfall']),
        'Fertilizer': float(data['fertilizer']),
        'Pesticide': float(data['pesticide'])
    }
    
    return input_data

def json_prediction_mode(input_source='stdin'):
    """Handle JSON input/output mode for predictions."""
    try:
        # Read JSON input
        if input_source == 'stdin':
            input_data = json.load(sys.stdin)
        else:
            with open(input_source, 'r') as f:
                input_data = json.load(f)
        
        # Validate and normalize input
        validated_data = validate_json_input(input_data)
        
        # Initialize predictor in quiet mode
        predictor = CropYieldPredictor(quiet=True)
        
        # Make prediction
        prediction, _ = predictor.predict_yield(validated_data)
        
        if isinstance(prediction, str) and 'Error' in prediction:
            error_result = {"error": prediction}
            print(json.dumps(error_result, indent=2))
            sys.exit(1)
        
        # Determine assessment
        if prediction > 3000:
            assessment = "Excellent yield expected"
        elif prediction > 2000:
            assessment = "Good yield expected"
        elif prediction > 1000:
            assessment = "Moderate yield expected"
        else:
            assessment = "Low yield expected"
        
        # Format and output JSON result
        result = format_json_output(prediction, validated_data['Area'], assessment)
        print(json.dumps(result, indent=2))
        
    except json.JSONDecodeError as e:
        error_result = {"error": f"Invalid JSON input: {str(e)}"}
        print(json.dumps(error_result, indent=2))
        sys.exit(1)
    except ValueError as e:
        error_result = {"error": str(e)}
        print(json.dumps(error_result, indent=2))
        sys.exit(1)
    except Exception as e:
        error_result = {"error": f"Prediction failed: {str(e)}"}
        print(json.dumps(error_result, indent=2))
        sys.exit(1)

def batch_prediction_from_csv(csv_file, output_file=None):
    """Make predictions for multiple records from a CSV file."""
    
    print(f"📄 Loading data from {csv_file}...")
    
    try:
        # Initialize predictor
        predictor = CropYieldPredictor()
        
        # Load CSV
        df = pd.read_csv(csv_file)
        print(f"📊 Loaded {len(df)} records for prediction.")
        
        # Make predictions
        results = []
        
        for idx, row in df.iterrows():
            print(f"🔄 Processing record {idx + 1}/{len(df)}...")
            
            prediction, _ = predictor.predict_yield(row.to_dict())
            
            result = row.to_dict()
            result['Predicted_Yield_RandomForest'] = prediction
            
            results.append(result)
        
        # Save results
        results_df = pd.DataFrame(results)
        
        if output_file is None:
            output_file = f"predictions_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
        
        results_df.to_csv(output_file, index=False)
        print(f"✅ Results saved to {output_file}")
        
        return results_df
        
    except Exception as e:
        print(f"❌ Error during batch prediction: {e}")
        return None

def main():
    """Main function to run the prediction system."""
    
    import sys
    import argparse
    
    parser = argparse.ArgumentParser(description="Crop Yield Predictor (CLI)")
    subparsers = parser.add_subparsers(dest="mode")
    
    # Batch mode
    batch_parser = subparsers.add_parser("batch", help="Batch prediction from CSV")
    batch_parser.add_argument("csv_file", help="Input CSV file with records to predict")
    batch_parser.add_argument("--out", dest="output_file", default=None, help="Output CSV file")
    
    # JSON mode
    json_parser = subparsers.add_parser("json", help="JSON input/output mode")
    json_parser.add_argument("--input", "-i", default="stdin", help="JSON input file (default: read from stdin)")
    
    # One-shot CLI mode
    one_parser = subparsers.add_parser("predict", help="One-shot prediction with CLI flags")
    one_parser.add_argument("--year", type=int, required=True, help="Crop year (e.g., 2024)")
    one_parser.add_argument("--state", type=str, required=True, help="State name")
    one_parser.add_argument("--crop", type=str, required=True, help="Crop name (e.g., Rice)")
    one_parser.add_argument("--season", type=str, required=True, help="Season (Kharif/Rabi/Summer/Whole Year/Autumn/Winter/Total)")
    one_parser.add_argument("--area", type=float, required=True, help="Area in hectares")
    one_parser.add_argument("--production", type=float, required=True, help="Production in tons")
    one_parser.add_argument("--rainfall", type=float, default=1000.0, help="Annual rainfall in mm (default 1000)")
    one_parser.add_argument("--fertilizer", type=float, default=50.0, help="Fertilizer usage in kg (default 50)")
    one_parser.add_argument("--pesticide", type=float, default=5.0, help="Pesticide usage in kg (default 5)")
    one_parser.add_argument("--district", type=str, default="Unknown", help="District name (optional)")
    
    # No args -> interactive
    args = parser.parse_args()
    
    if args.mode == "batch":
        batch_prediction_from_csv(args.csv_file, args.output_file)
        return
    
    if args.mode == "json":
        json_prediction_mode(args.input)
        return
    
    if args.mode == "predict":
        # Build input data dict from args
        input_data = {
            'Crop_Year': args.year,
            'State': args.state,
            'District': args.district,
            'Crop': args.crop,
            'Season': args.season,
            'Area': args.area,
            'Production': args.production,
            'Annual_Rainfall': args.rainfall,
            'Fertilizer': args.fertilizer,
            'Pesticide': args.pesticide,
        }
        
        # Run prediction
        predictor = CropYieldPredictor()
        prediction, _ = predictor.predict_yield(input_data)
        
        if isinstance(prediction, str) and 'Error' in prediction:
            print(f"Error: {prediction}")
            sys.exit(1)
        
        print("Prediction result:")
        print(f"Random Forest: {prediction:.2f} kg/hectare")
        
        total_prod = (prediction * args.area) / 1000.0
        print(f"Total expected production: {total_prod:.2f} tons")
        
        # Interpretation
        if prediction > 3000:
            print("Assessment: 🟢 Excellent yield expected!")
        elif prediction > 2000:
            print("Assessment: 🟡 Good yield expected.")
        elif prediction > 1000:
            print("Assessment: 🟠 Moderate yield expected.")
        else:
            print("Assessment: 🔴 Low yield expected.")
        return
    
    # Default interactive mode
    interactive_prediction()

if __name__ == "__main__":
    main()