import matplotlib.pyplot as plt
import numpy as np
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
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import make_pipeline
import os

def load_petal_data():
    """
    Loads the Iris dataset and extracts petal dimensions.
    """
    iris = datasets.load_iris(as_frame=True)
    df = iris.frame
    X = df[['petal length (cm)']]
    y = df['petal width (cm)']
    return X, y

def train_and_evaluate_poly_model(degree, X_train, y_train, X_test, y_test):
    """
    Trains a polynomial regression model of a given degree.
    """
    # Create a pipeline that first transforms features, then applies linear regression
    model = make_pipeline(PolynomialFeatures(degree), linear_model.LinearRegression())
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)

    print(f"Degree {degree} Model:")
    print(f"  MSE: {mse:.4f}")
    print(f"  R^2: {r2:.4f}")

    return model

def visualize_polynomial_fits(models, X, y):
    """
    Plots the data and the regression curves for different degrees.
    """
    if not os.path.exists('plots'):
        os.makedirs('plots')
        
    plt.figure(figsize=(10, 6))
    plt.scatter(X, y, color='black', alpha=0.5, label='Actual Data')
    
    # Generate a smooth range of X values for plotting curves
    X_plot = np.linspace(X.min(), X.max(), 100).reshape(-1, 1)
    
    colors = ['blue', 'green', 'red']
    
    for degree, model in models.items():
        y_plot = model.predict(X_plot)
        plt.plot(X_plot, y_plot, color=colors[degree-1], linewidth=2, label=f'Degree {degree}')
        
    plt.xlabel('Petal Length (cm)')
    plt.ylabel('Petal Width (cm)')
    plt.title('Polynomial Regression Comparison')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    output_path = 'plots/polynomial_regression.png'
    plt.savefig(output_path)
    plt.close()
    print(f"Polynomial plot saved to {output_path}")

def main():
    print("Starting Polynomial Regression Analysis...")
    
    X, y = load_petal_data()
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    models = {}
    degrees = [1, 2, 3]
    
    for degree in degrees:
        models[degree] = train_and_evaluate_poly_model(degree, X_train, y_train, X_test, y_test)
        
    visualize_polynomial_fits(models, X, y)
    print("Analysis Complete.")

if __name__ == "__main__":
    main()