import matplotlib.pyplot as plt
import seaborn as sns
from sklearn import datasets
import pandas as pd
import os

def load_data():
    """Loads the Iris dataset into a Pandas DataFrame."""
    iris = datasets.load_iris(as_frame=True)
    df = iris.frame
    # Map target integers to target names for better labeling
    df['species'] = df['target'].map(lambda x: iris.target_names[x])
    # Drop the numeric target column as we have the species name now
    df = df.drop(columns=['target'])
    return df

def setup_plots_dir():
    """Creates the plots directory if it doesn't exist."""
    if not os.path.exists('plots'):
        os.makedirs('plots')

def plot_pairplot(df):
    """Generates and saves a pairplot."""
    plt.figure(figsize=(10, 8))
    sns.pairplot(df, hue='species')
    plt.savefig('plots/pairplot.png')
    plt.close()
    print("Saved plots/pairplot.png")

def plot_correlation_heatmap(df):
    """Generates and saves a correlation heatmap."""
    plt.figure(figsize=(8, 6))
    # Select only numeric columns for correlation
    numeric_df = df.select_dtypes(include=['float64', 'int64'])
    corr = numeric_df.corr()
    sns.heatmap(corr, annot=True, cmap='coolwarm', fmt=".2f")
    plt.title('Feature Correlation Heatmap')
    plt.tight_layout()
    plt.savefig('plots/heatmap.png')
    plt.close()
    print("Saved plots/heatmap.png")

def plot_histograms(df):
    """Generates and saves histograms for each feature."""
    df.hist(figsize=(10, 8), bins=20)
    plt.suptitle('Feature Distributions')
    plt.tight_layout()
    plt.savefig('plots/histograms.png')
    plt.close()
    print("Saved plots/histograms.png")

def main():
    setup_plots_dir()
    df = load_data()
    
    print("Generating plots...")
    plot_pairplot(df)
    plot_correlation_heatmap(df)
    plot_histograms(df)
    print("Done!")

if __name__ == "__main__":
    main()