import pandas as pd
import numpy as np
import joblib

from pathlib import Path

from sklearn.model_selection import (
    train_test_split,
    RandomizedSearchCV,
    KFold
)

from sklearn.ensemble import RandomForestRegressor

from sklearn.metrics import (
    mean_absolute_error,
    mean_squared_error,
    r2_score
)


# ============================================================
# 1. File paths
# ============================================================

# Get the folder where this Python file is located
BASE_DIR = Path(__file__).parent

# Input cleaned dataset
INPUT_FILE = BASE_DIR / "cleaned_data.csv"

# Output file for the tuned model
MODEL_FILE = BASE_DIR / "tuned_random_forest.pkl"

# Output file containing the tuning results
RESULT_FILE = BASE_DIR / "random_forest_tuning_results.csv"


# ============================================================
# 2. Load cleaned dataset
# ============================================================

print("Loading cleaned dataset...")

df = pd.read_csv(INPUT_FILE)

print("Dataset shape:", df.shape)


# ============================================================
# 3. Define features and target
# ============================================================

# These are the same features used in our previous models.
features = [
    "cylinders",
    "displacement",
    "horsepower",
    "weight",
    "acceleration",
    "model year",
    "origin"
]

# Target variable: vehicle fuel efficiency in MPG
target = "mpg"

X = df[features]
y = df[target]


# ============================================================
# 4. Create train/test split
# ============================================================

# IMPORTANT:
# The test set is kept separate and is NOT used during
# hyperparameter tuning.
#
# This gives us an unbiased final evaluation of the tuned model.

X_train, X_test, y_train, y_test = train_test_split(
    X,
    y,
    test_size=0.20,
    random_state=42
)

print("\nTraining samples:", len(X_train))
print("Testing samples :", len(X_test))


# ============================================================
# 5. Create the Random Forest model
# ============================================================

rf = RandomForestRegressor(
    random_state=42,
    n_jobs=-1
)


# ============================================================
# 6. Define hyperparameter search space
# ============================================================

# Hyperparameters control how the Random Forest creates
# and combines its decision trees.
#
# We are testing multiple values instead of manually choosing
# a single configuration.

param_distributions = {

    # Number of trees in the forest
    "n_estimators": [
        100,
        200,
        300,
        500,
        800
    ],

    # Maximum depth of each decision tree
    # None means the tree can grow until stopping criteria
    "max_depth": [
        None,
        5,
        10,
        15,
        20,
        30
    ],

    # Minimum number of samples required to split a node
    "min_samples_split": [
        2,
        5,
        10,
        15
    ],

    # Minimum number of samples required at a leaf node
    "min_samples_leaf": [
        1,
        2,
        4,
        8
    ],

    # Number of features considered when looking for a split
    "max_features": [
        1.0,
        "sqrt",
        "log2"
    ]
}


# ============================================================
# 7. 5-Fold Cross Validation
# ============================================================

# The training data is divided into five parts.
# Four parts are used for training and one part for validation.
# This process is repeated five times.
#
# This helps us select hyperparameters that generalize better
# instead of relying on one particular validation split.

cv = KFold(
    n_splits=5,
    shuffle=True,
    random_state=42
)


# ============================================================
# 8. Randomized Hyperparameter Search
# ============================================================

# RandomizedSearchCV tests a selected number of random
# combinations from the parameter space.
#
# It is generally faster than testing every possible
# combination with GridSearchCV.

search = RandomizedSearchCV(
    estimator=rf,
    param_distributions=param_distributions,

    # Number of random combinations to test
    n_iter=50,

    # Use 5-fold cross-validation
    cv=cv,

    # R² is used to select the best model
    scoring="r2",

    # Use all available CPU cores
    n_jobs=-1,

    # Store training scores as well
    return_train_score=True,

    # Reproducible random search
    random_state=42,

    # Show progress in the terminal
    verbose=2
)


# ============================================================
# 9. Start hyperparameter tuning
# ============================================================

print("\n==============================================")
print("STARTING RANDOM FOREST HYPERPARAMETER TUNING")
print("==============================================")

print("\nTesting 50 random hyperparameter combinations...")
print("Using 5-fold cross-validation...")


search.fit(X_train, y_train)


# ============================================================
# 10. Display best parameters
# ============================================================

print("\n==============================================")
print("BEST HYPERPARAMETERS")
print("==============================================")

print("\nBest parameters found:")

for parameter, value in search.best_params_.items():
    print(f"{parameter}: {value}")


# ============================================================
# 11. Best cross-validation score
# ============================================================

print("\nBest Cross-Validation R²:")
print(f"{search.best_score_:.4f}")


# ============================================================
# 12. Get the best model
# ============================================================

best_model = search.best_estimator_


# ============================================================
# 13. Evaluate on untouched test set
# ============================================================

# The test set was never used during hyperparameter tuning.
# This is the most important final performance measurement.

y_pred = best_model.predict(X_test)


# Calculate MAE
mae = mean_absolute_error(
    y_test,
    y_pred
)


# Calculate RMSE
rmse = np.sqrt(
    mean_squared_error(
        y_test,
        y_pred
    )
)


# Calculate R²
r2 = r2_score(
    y_test,
    y_pred
)


print("\n==============================================")
print("TUNED RANDOM FOREST TEST PERFORMANCE")
print("==============================================")

print(f"R²   : {r2:.4f}")
print(f"RMSE : {rmse:.4f}")
print(f"MAE  : {mae:.4f}")


# ============================================================
# 14. Compare with previous Random Forest
# ============================================================

# Previous untuned Random Forest:
# R²   = 0.8923
# RMSE = 2.3443
# MAE  = 1.6481

previous_r2 = 0.8923
previous_rmse = 2.3443
previous_mae = 1.6481

print("\n==============================================")
print("COMPARISON WITH PREVIOUS RANDOM FOREST")
print("==============================================")

print("\nMetric          Previous       Tuned")

print(
    f"R²              {previous_r2:.4f}        {r2:.4f}"
)

print(
    f"RMSE             {previous_rmse:.4f}        {rmse:.4f}"
)

print(
    f"MAE              {previous_mae:.4f}        {mae:.4f}"
)


# ============================================================
# 15. Save all hyperparameter search results
# ============================================================

# This CSV allows us to inspect how different parameter
# combinations performed.

results_df = pd.DataFrame(
    search.cv_results_
)

results_df = results_df.sort_values(
    by="rank_test_score"
)

results_df.to_csv(
    RESULT_FILE,
    index=False
)

print("\nTuning results saved to:")
print(RESULT_FILE)


# ============================================================
# 16. Save tuned model
# ============================================================

joblib.dump(
    best_model,
    MODEL_FILE
)

print("\n==============================================")
print("TUNED MODEL SAVED")
print("==============================================")

print("Model file:")
print(MODEL_FILE)

print("\nHyperparameter tuning complete!")