import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
import joblib


def load_and_prepare_data():
    """
    Load the dataset and prepare it for modeling.
    
    Returns:
        tuple: X (features), y (target)
    """
    print("Loading data...")
    df = pd.read_csv("New_York_cars.csv")
    
    # Remove rows where target 'money' is missing
    df = df.dropna(subset=["money"])
    
    # Define features and target
    features = [
        "Year",
        "Mileage",
        "brand",
        "new&used",
        "Exterior color",
        "Interior color",
        "Drivetrain",
        "MPG",
        "Fuel type",
        "Transmission",
        "Engine",
        "Convenience",
        "Entertainment",
        "Exterior",
        "Safety",
        "Seating",
        "Accidents or damage",
        "Clean title",
        "1-owner vehicle",
        "Personal use only",
        "Model",
    ]
    target = "money"
    
    X = df[features]
    y = df[target]
    
    return X, y


def create_preprocessor():
    """
    Create a preprocessing pipeline for handling both numeric and categorical features.
    
    Returns:
        ColumnTransformer: Preprocessing pipeline
    """
    # Define numeric and categorical features
    numeric_features = ["Year", "Mileage"]
    categorical_features = [
        "brand",
        "MPG",
        "Fuel type",
        "Transmission",
        "Engine",
        "Safety",
        "Accidents or damage",
        "Clean title",
        "1-owner vehicle",
        "Personal use only",
        "Model",
    ]
    
    # Numeric preprocessing pipeline
    numeric_transformer = Pipeline(
        steps=[("imputer", SimpleImputer(strategy="median"))]
    )
    
    # Categorical preprocessing pipeline
    categorical_transformer = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="constant", fill_value="Unknown")),
            ("onehot", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
        ]
    )
    
    # Combine preprocessors
    preprocessor = ColumnTransformer(
        transformers=[
            ("num", numeric_transformer, numeric_features),
            ("cat", categorical_transformer, categorical_features),
        ]
    )
    
    return preprocessor


def train_model(X_train, y_train, preprocessor):
    """
    Train a linear regression model with preprocessing.
    
    Args:
        X_train: Training features
        y_train: Training target
        preprocessor: Preprocessing pipeline
        
    Returns:
        Pipeline: Trained model pipeline
    """
    print("Training Linear Regression...")
    
    # Create model pipeline
    lr_pipeline = Pipeline(
        steps=[("preprocessor", preprocessor), ("regressor", LinearRegression())]
    )
    
    # Fit the model
    lr_pipeline.fit(X_train, y_train)
    
    return lr_pipeline


def evaluate_model(model, X_test, y_test):
    """
    Evaluate the trained model using MSE and R2 metrics.
    
    Args:
        model: Trained model pipeline
        X_test: Test features
        y_test: Test target
        
    Returns:
        tuple: (mse, r2)
    """
    # Make predictions
    y_pred = model.predict(X_test)
    
    # Calculate metrics
    mse = mean_squared_error(y_test, y_pred)
    rmse = np.sqrt(mse)
    r2 = r2_score(y_test, y_pred)
    
    print(f"Linear Regression MSE: {mse}")
    print(f"Linear Regression RMSE: {rmse:.4f}")
    print(f"Linear Regression R2: {r2}")
    
    return mse, r2


def save_model(model, filename="car_price_model.joblib"):
    """
    Save the trained model to disk.
    
    Args:
        model: Trained model pipeline
        filename: Output filename
    """
    joblib.dump(model, filename)
    print(f"Model saved to {filename}")


def train():
    """
    Main training function that orchestrates the entire process.
    """
    # Load and prepare data
    X, y = load_and_prepare_data()
    
    # Split data into train and test sets
    print("Splitting data...")
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    
    # Create preprocessing pipeline
    preprocessor = create_preprocessor()
    
    # Train the model
    model = train_model(X_train, y_train, preprocessor)
    
    # Evaluate the model
    mse, r2 = evaluate_model(model, X_test, y_test)
    
    # Save the model
    save_model(model)


if __name__ == "__main__":
    train()