import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestRegressor

def get_data():
    # Load data
    df = pd.read_csv('household_power_consumption.txt', sep=';', 
        na_values=['?']
    )

    # Preprocess
    df = df.dropna()
    df['datetime'] = pd.to_datetime(
        df['Date'] + ' ' + df['Time'],
        format='%d/%m/%Y %H:%M:%S'
    )
    df = df.drop(['Date', 'Time'], axis=1)
    df['hour'] = df['datetime'].dt.hour
    df['day'] = df['datetime'].dt.dayofweek
    df['month'] = df['datetime'].dt.month

    return df

def analyze_relationships(df):
    # Calculate correlations
    correlations = df[['Global_intensity', 'Global_active_power', 
        'Global_reactive_power', 'Voltage']
    ].corr()

    # Plot correlation matrix
    plt.figure(figsize=(10, 8))
    sns.heatmap(correlations, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
    plt.title('Correlation Matrix of Power Measurements')
    plt.show()

    # Scatter plot of Active Power vs Intensity
    plt.figure(figsize=(10, 6))
    plt.scatter(df['Global_active_power'], df['Global_intensity'], alpha=0.1)
    plt.xlabel('Global Active Power (kW)')
    plt.ylabel('Global Intensity (A)')
    plt.title('Active Power vs Intensity')
    plt.show()

def train_model(X, y):
    # Split data
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    # Scale features
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    models = {
        'Linear Regression': LinearRegression(),
        'Random Forest': RandomForestRegressor(n_estimators=100, random_state=42)
    }

    results = {}
    for name, model in models.items():
        # Train model
        model.fit(X_train_scaled, y_train)

        # Make predictions
        y_pred = model.predict(X_test_scaled)

        # Calculate metrics
        mse = mean_squared_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)

        results[name] = {
            'model': model,
            'mse': mse,
            'r2': r2,
            'predictions': y_pred,
            'actual': y_test
        }

    return results, scaler, X_test, y_test

def plot_results(results):
    plt.figure(figsize=(15, 10))

    for i, (name, result) in enumerate(results.items(), 1):
        plt.subplot(2, 1, i)
        plt.scatter(result['actual'], result['predictions'], alpha=0.5)
        plt.plot(
            [result['actual'].min(), result['actual'].max()], 
            [result['actual'].min(), result['actual'].max()], 
            'r--', lw=2
        )
        plt.xlabel('Actual Energy Consumption (Global_intensity)')
        plt.ylabel('Predicted Energy Consumption')
        plt.title(f'{name}\nR² Score: {result["r2"]:.3f}, MSE: {result["mse"]:.3f}')

    plt.tight_layout()
    plt.show()

def feature_importance(results, X):
    # For Random Forest model
    model = results['Random Forest']['model']
    feature_imp = pd.DataFrame({
        'feature': X.columns,
        'importance': model.feature_importances_
    }).sort_values('importance', ascending=False)

    plt.figure(figsize=(10, 6))
    sns.barplot(x='importance', y='feature', data=feature_imp)
    plt.title('Feature Importance in Energy Consumption Prediction')
    plt.show()

    return feature_imp

def main():
    df = get_data()

    # Optional for feature selection
    # analyze_relationships(df)

    # Set features and target
    features = [
        'hour', 'day', 'month', 
        'Sub_metering_1', 'Sub_metering_2', 'Sub_metering_3'
    ]

    X = df[features]
    y = df['Global_intensity']

    results, scaler, X_test, y_test = train_model(X, y)
    plot_results(results)

    # Analyze feature importance
    df = feature_importance(results, X)
    print("\nFeature Importance:")
    print(df)

    # Example prediction
    print("\nExample Prediction:")
    sample_data = X_test[:1]
    sample_scaled = scaler.transform(sample_data)

    for name, result in results.items():
        prediction = result['model'].predict(sample_scaled)[0]
        print(f"{name} prediction: {prediction:.2f} kW")
        print(f"Actual value: {y_test.iloc[0]:.2f} kW")

if __name__ == "__main__":
    main()