import matplotlib.pyplot as plt
import pandas as pd
from sklearn import datasets, linear_model
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
import os

def load_petal_data():
    """
    Loads the Iris dataset and extracts petal dimensions.
    
    Returns:
        X (DataFrame): Feature matrix containing 'petal length (cm)'.
        y (Series): Target vector containing 'petal width (cm)'.
    """
    iris = datasets.load_iris(as_frame=True)
    df = iris.frame
    
    # We want to predict Petal Width based on Petal Length
    # Reshaping X to be 2D array as required by sklearn (n_samples, n_features)
    X = df[['petal length (cm)']]
    y = df['petal width (cm)']
    
    return X, y

def train_linear_model(X_train, y_train):
    """
    Trains a Linear Regression model.
    
    Args:
        X_train: Training features.
        y_train: Training target.
        
    Returns:
        model: Trained LinearRegression model.
    """
    model = linear_model.LinearRegression()
    model.fit(X_train, y_train)
    return model

def evaluate_model_performance(model, X_test, y_test):
    """
    Evaluates the model using Mean Squared Error and R-squared score.
    
    Args:
        model: Trained model.
        X_test: Test features.
        y_test: Test target.
    """
    y_pred = model.predict(X_test)
    
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    
    print("Model Performance:")
    print(f"  Coefficients: {model.coef_[0]:.4f}")
    print(f"  Intercept: {model.intercept_:.4f}")
    print(f"  Mean Squared Error (MSE): {mse:.4f}")
    print(f"  Coefficient of Determination (R^2): {r2:.4f}")

def visualize_regression_line(model, X, y):
    """
    Plots the data points and the regression line.
    
    Args:
        model: Trained model.
        X: All features (for plotting).
        y: All targets (for plotting).
    """
    if not os.path.exists('plots'):
        os.makedirs('plots')
        
    plt.figure(figsize=(10, 6))
    
    # Plot actual data points
    plt.scatter(X, y, color='black', label='Actual Data')
    
    # Plot regression line
    # We predict on the full X range to draw the line
    y_pred_full = model.predict(X)
    plt.plot(X, y_pred_full, color='blue', linewidth=3, label='Regression Line')
    
    plt.xlabel('Petal Length (cm)')
    plt.ylabel('Petal Width (cm)')
    plt.title('Linear Regression: Petal Length vs Petal Width')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    output_path = 'plots/linear_regression.png'
    plt.savefig(output_path)
    plt.close()
    print(f"Regression plot saved to {output_path}")

def main():
    print("Starting Linear Regression Analysis...")
    
    # 1. Load Data
    X, y = load_petal_data()
    
    # 2. Split Data, 80/20 rule
    X_train, X_test, y_train, y_test = train_test_split(X, y)#, test_size=0.2) 
    
    # 3. Train Model
    model = train_linear_model(X_train, y_train)
    
    # 4. Evaluate
    evaluate_model_performance(model, X_test, y_test)
    
    # 5. Visualize
    visualize_regression_line(model, X, y)
    
    print("Analysis Complete.")

if __name__ == "__main__":
    main()