"""
Monthly Budget Forecasting with Real Dataset
---------------------------------------------
End-to-end pipeline:
1. Load real dataset from CSV (monthly_forecast_dataset_large.csv)
2. Clean and preprocess data
3. Engineer temporal features
4. Train multiple regression models
5. Compare performance and select best
6. Evaluate on test set
7. Forecast next 3 months
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
from pathlib import Path
import joblib

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import Ridge, Lasso, LinearRegression
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

import warnings
warnings.filterwarnings('ignore')


# ============================================================================
# PART 1: LOAD REAL DATASET FROM CSV
# ============================================================================

def load_real_budget_data(csv_path):
    """
    Load real budget dataset from CSV.
    
    Args:
        csv_path: Path to the CSV file
    
    Returns:
        pandas DataFrame with budget data
    """
    print("="*70)
    print("LOADING REAL DATASET FROM CSV")
    print("="*70)
    
    df = pd.read_csv(csv_path)
    # Restrict to month and monthly_budget_pkr for single-family series
    keep_cols = ['month', 'monthly_budget_pkr']
    df = df[keep_cols]
    
    print(f"\n✓ Loaded {len(df)} rows from {csv_path}")
    print(f"\nDataset shape: {df.shape}")
    print(f"\nColumns: {list(df.columns)}")
    print(f"\nFirst 10 rows:")
    print(df.head(10))
    print(f"\nStatistics:")
    print(df.describe().round(2))
    
    return df


# ============================================================================
# PART 2: DATA CLEANING & PREPROCESSING
# ============================================================================

def clean_and_preprocess(df):
    """
    Clean data and engineer features for the real dataset.
    
    Args:
        df: Input DataFrame
    
    Returns:
        Cleaned DataFrame with engineered features
    """
    print("\n" + "="*70)
    print("DATA CLEANING & PREPROCESSING")
    print("="*70)
    
    df = df.copy()
    # Keep only month and target for single-family series
    df = df[['month', 'monthly_budget_pkr']]

    # Handle missing values
    missing = df.isnull().sum()
    if missing.any():
        print(f"\n⚠ Missing values detected:\n{missing[missing > 0]}")
        df = df.dropna()

    # Temporal features
    df['month_date'] = pd.to_datetime(df['month'])
    df['time_idx'] = (df['month_date'] - df['month_date'].min()).dt.days // 30
    df['month_num'] = df['month_date'].dt.month
    df['month_sin'] = np.sin(2 * np.pi * df['month_num'] / 12)
    df['month_cos'] = np.cos(2 * np.pi * df['month_num'] / 12)

    print(f"\n✓ Data cleaned successfully")
    print(f"  Final shape: {df.shape}")
    print(f"  Missing values: {df.isnull().sum().sum()}")
    print(f"  Columns: {list(df.columns)}")
    
    return df


# ============================================================================
# PART 3: MODEL TRAINING & COMPARISON
# ============================================================================

def train_and_compare_models(X_train, y_train, X_test, y_test):
    """
    Train multiple models and compare performance.
    
    Args:
        X_train, y_train: Training data
        X_test, y_test: Test data
    
    Returns:
        Dictionary with model results
    """
    print("\n" + "="*70)
    print("TRAINING & COMPARING MULTIPLE MODELS")
    print("="*70)
    
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    models = {
        'Linear Regression': LinearRegression(),
        'Ridge Regression': Ridge(alpha=1.0, random_state=42),
        'Lasso Regression': Lasso(alpha=0.1, random_state=42),
    }
    
    results = {}
    
    for model_name, model in models.items():
        print(f"\n--- Training {model_name} ---")
        
        # Train
        model.fit(X_train_scaled, y_train)
        
        # Predict
        y_pred = model.predict(X_test_scaled)
        
        # Evaluate
        mae = mean_absolute_error(y_test, y_pred)
        rmse = np.sqrt(mean_squared_error(y_test, y_pred))
        r2 = r2_score(y_test, y_pred)
        mape = np.mean(np.abs((y_test - y_pred) / (y_test + 1e-6))) * 100
        
        results[model_name] = {
            'model': model,
            'scaler': scaler,
            'mae': mae,
            'rmse': rmse,
            'r2': r2,
            'mape': mape,
            'predictions': y_pred
        }
        
        print(f"  MAE:  {mae:>10.2f}")
        print(f"  RMSE: {rmse:>10.2f}")
        print(f"  R²:   {r2:>10.4f}")
        print(f"  MAPE: {mape:>10.2f}%")
    
    # Find best model
    best_model_name = max(results, key=lambda x: results[x]['r2'])
    print(f"\n" + "="*70)
    print(f"🏆 BEST MODEL: {best_model_name}")
    print(f"   R² Score: {results[best_model_name]['r2']:.4f}")
    print("="*70)
    
    return results, best_model_name, scaler


# ============================================================================
# PART 4: EVALUATION & VISUALIZATION
# ============================================================================

def evaluate_and_visualize(results, best_model_name, X_test, y_test, output_dir='output'):
    """
    Detailed evaluation and visualization of best model.
    """
    print("\n" + "="*70)
    print("DETAILED EVALUATION OF BEST MODEL")
    print("="*70)
    
    output_path = Path(output_dir)
    output_path.mkdir(exist_ok=True)
    
    best_result = results[best_model_name]
    y_pred = best_result['predictions']
    
    # Performance metrics table
    print("\nCOMPARATIVE MODEL PERFORMANCE:")
    print("-" * 70)
    print(f"{'Model':<20} {'MAE':>12} {'RMSE':>12} {'R²':>12} {'MAPE':>10}")
    print("-" * 70)
    for model_name in sorted(results.keys()):
        metrics = results[model_name]
        print(f"{model_name:<20} {metrics['mae']:>12.2f} {metrics['rmse']:>12.2f} "
              f"{metrics['r2']:>12.4f} {metrics['mape']:>10.2f}%")
    print("-" * 70)
    
    # Visualization
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    # 1. Actual vs Predicted
    ax = axes[0, 0]
    ax.scatter(y_test, y_pred, alpha=0.6, s=50)
    ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
    ax.set_xlabel('Actual Budget (PKR)', fontsize=11)
    ax.set_ylabel('Predicted Budget (PKR)', fontsize=11)
    ax.set_title(f'{best_model_name} - Actual vs Predicted', fontsize=12, fontweight='bold')
    ax.grid(True, alpha=0.3)
    
    # 2. Residuals
    ax = axes[0, 1]
    residuals = y_test - y_pred
    ax.scatter(y_pred, residuals, alpha=0.6, s=50)
    ax.axhline(y=0, color='r', linestyle='--', lw=2)
    ax.set_xlabel('Predicted Budget (PKR)', fontsize=11)
    ax.set_ylabel('Residuals (PKR)', fontsize=11)
    ax.set_title('Residual Plot', fontsize=12, fontweight='bold')
    ax.grid(True, alpha=0.3)
    
    # 3. Distribution of Residuals
    ax = axes[1, 0]
    ax.hist(residuals, bins=20, edgecolor='black', alpha=0.7)
    ax.set_xlabel('Residuals (PKR)', fontsize=11)
    ax.set_ylabel('Frequency', fontsize=11)
    ax.set_title('Distribution of Residuals', fontsize=12, fontweight='bold')
    ax.axvline(x=0, color='r', linestyle='--', lw=2)
    ax.grid(True, alpha=0.3, axis='y')
    
    # 4. Model Comparison
    ax = axes[1, 1]
    model_names = list(results.keys())
    r2_scores = [results[m]['r2'] for m in model_names]
    colors = ['green' if m == best_model_name else 'lightblue' for m in model_names]
    ax.barh(model_names, r2_scores, color=colors, edgecolor='black')
    ax.set_xlabel('R² Score', fontsize=11)
    ax.set_title('Model Comparison (R² Score)', fontsize=12, fontweight='bold')
    ax.set_xlim([min(0, min(r2_scores) - 0.1), 1])
    for i, v in enumerate(r2_scores):
        ax.text(v + 0.01, i, f'{v:.4f}', va='center', fontsize=10)
    
    plt.tight_layout()
    plot_path = output_path / f'{best_model_name.replace(" ", "_")}_evaluation.png'
    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
    print(f"\n✓ Evaluation plot saved to {plot_path}")
    plt.close()
    
    # Save metrics to CSV
    metrics_df = pd.DataFrame(results).T[['mae', 'rmse', 'r2', 'mape']]
    metrics_path = output_path / 'model_comparison.csv'
    metrics_df.to_csv(metrics_path)
    print(f"✓ Metrics comparison saved to {metrics_path}")


# ============================================================================
# PART 5: FUTURE PREDICTIONS
# ============================================================================

def predict_future_budgets(df, best_result, feature_columns, n_future_months=3):
    """
    Predict budget for next n months.
    
    Args:
        df: Full preprocessed dataframe
        best_result: Best model result dictionary
        feature_columns: List of feature column names
        n_future_months: Number of months to forecast
    
    Returns:
        DataFrame with future predictions
    """
    print("\n" + "="*70)
    print(f"FORECASTING NEXT {n_future_months} MONTHS")
    print("="*70)
    
    model = best_result['model']
    scaler = best_result['scaler']

    last_month = df['month_date'].max()
    future_dates = [last_month + timedelta(days=30 * (i + 1)) for i in range(n_future_months)]
    future_months = [d.strftime('%Y-%m') for d in future_dates]

    base_row = df.iloc[-1].copy()
    future_predictions = []

    for future_date, future_month in zip(future_dates, future_months):
        future_row = base_row.copy()
        future_row['month'] = future_month
        future_row['month_date'] = future_date
        future_row['time_idx'] = (future_date - df['month_date'].min()).days // 30
        future_row['month_num'] = future_date.month
        future_row['month_sin'] = np.sin(2 * np.pi * future_row['month_num'] / 12)
        future_row['month_cos'] = np.cos(2 * np.pi * future_row['month_num'] / 12)

        X_future = future_row[feature_columns].values.reshape(1, -1)
        X_future_scaled = scaler.transform(X_future)
        predicted_budget = model.predict(X_future_scaled)[0]

        future_predictions.append({
            'month': future_month,
            'predicted_monthly_budget_pkr': predicted_budget,
        })

    future_df = pd.DataFrame(future_predictions)
    
    print("\nFuture Budget Forecasts:")
    print("-" * 70)
    print(future_df.to_string(index=False))
    print("-" * 70)
    
    return future_df


# ============================================================================
# PART 6: MODEL PERSISTENCE & RELOAD PREDICTION
# ============================================================================

def save_best_model(best_result, best_model_name, feature_columns, output_dir):
    """Persist best model, scaler, and feature list to disk."""
    output_path = Path(output_dir)
    output_path.mkdir(exist_ok=True)
    model_bundle = {
        'model': best_result['model'],
        'scaler': best_result['scaler'],
        'feature_columns': feature_columns,
        'model_name': best_model_name,
    }
    model_path = output_path / f"best_model_{best_model_name.replace(' ', '_').lower()}.joblib"
    joblib.dump(model_bundle, model_path)
    print(f"\n✓ Saved best model bundle to {model_path}")
    return model_path


def load_model(model_path):
    """Load persisted model bundle from disk."""
    model_bundle = joblib.load(model_path)
    print(f"✓ Loaded model bundle from {model_path}")
    return model_bundle


def predict_with_loaded_model(df, model_bundle, n_future_months=3):
    """Forecast using a reloaded model bundle."""
    best_result = {
        'model': model_bundle['model'],
        'scaler': model_bundle['scaler'],
    }
    feature_columns = model_bundle['feature_columns']
    return predict_future_budgets(df, best_result, feature_columns, n_future_months)


# ============================================================================
# PART 7: MAIN EXECUTION
# ============================================================================

def main():
    """Main execution pipeline"""
    print("\n")
    print("█" * 70)
    print("█" + " " * 68 + "█")
    print("█" + "  BUDGET FORECASTING WITH REAL DATASET - PIPELINE".center(68) + "█")
    print("█" + " " * 68 + "█")
    print("█" * 70)

    csv_path = 'd:/FoodData/Monthly forcast/monthly_budget_single_family_24m.csv'
    output_dir = 'output'
    Path(output_dir).mkdir(exist_ok=True)

    # Step 1: Load real dataset
    df = load_real_budget_data(csv_path)

    # Step 2: Clean and preprocess
    df_processed = clean_and_preprocess(df)

    # Step 3: Prepare features
    feature_cols = ['time_idx', 'month_num', 'month_sin', 'month_cos']
    X = df_processed[feature_cols]
    y = df_processed['monthly_budget_pkr']

    # Step 4: Train-test split (chronological)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, shuffle=False
    )

    print(f"\n" + "="*70)
    print("DATA SPLIT")
    print("="*70)
    print(f"Training set:   {len(X_train)} samples")
    print(f"Test set:       {len(X_test)} samples")
    print(f"Features:       {X_train.shape[1]}")
    print(f"\nFeature columns used ({len(feature_cols)}):")
    for i, col in enumerate(feature_cols, 1):
        print(f"  {i:2d}. {col}")

    # Step 5: Train and compare models
    results, best_model_name, scaler = train_and_compare_models(
        X_train, y_train, X_test, y_test
    )

    # Step 6: Evaluate and visualize
    evaluate_and_visualize(results, best_model_name, X_test, y_test, output_dir)

    # Step 7: Make future predictions (next month by default)
    best_result = results[best_model_name]
    future_df = predict_future_budgets(
        df_processed, best_result, feature_cols, n_future_months=1
    )

    # Save future predictions
    future_path = Path(output_dir) / 'future_predictions_real_data.csv'
    future_df.to_csv(future_path, index=False)
    print(f"\n✓ Future prediction saved to {future_path}")

    # Save and reload the best model, then predict again to validate persistence
    model_path = save_best_model(best_result, best_model_name, feature_cols, output_dir)
    loaded_bundle = load_model(model_path)
    future_df_loaded = predict_with_loaded_model(
        df_processed, loaded_bundle, n_future_months=1
    )
    future_reloaded_path = Path(output_dir) / 'future_predictions_real_data_reloaded.csv'
    future_df_loaded.to_csv(future_reloaded_path, index=False)
    print(f"✓ Future prediction (reloaded model) saved to {future_reloaded_path}")

    # Save best model info
    model_info_path = Path(output_dir) / 'best_model_info.txt'
    with open(model_info_path, 'w') as f:
        f.write(f"BEST MODEL: {best_model_name}\n")
        f.write(f"R² Score: {results[best_model_name]['r2']:.4f}\n")
        f.write(f"MAE: {results[best_model_name]['mae']:.2f}\n")
        f.write(f"RMSE: {results[best_model_name]['rmse']:.2f}\n")
        f.write(f"MAPE: {results[best_model_name]['mape']:.2f}%\n")
        f.write(f"\nFeatures used: {len(feature_cols)}\n")
        for col in feature_cols:
            f.write(f"  - {col}\n")

    print(f"✓ Model info saved to {model_info_path}")

    print("\n" + "█" * 70)
    print("█" + "  ✓ PIPELINE EXECUTION COMPLETED SUCCESSFULLY".center(68) + "█")
    print("█" * 70 + "\n")


if __name__ == '__main__':
    main()