# ----------------------------------------------------------------------------
# Custom functions -----------------------------------------------------------
# ----------------------------------------------------------------------------

import numpy as np
from scipy.linalg import block_diag

def make_group_data(
        n_obs: int, n_groups: int, group_size: int, corr: float, seed: int,
        target_irreducible_mse: float = 200.0, snr: float = 4.0):
    """
    Generates grouped-feature data with a controlled signal-to-noise ratio.
    """
    np.random.seed(seed)

    n_features = n_groups * group_size
    group_indices = np.repeat(np.arange(n_groups), group_size)
    beta = np.zeros(n_features)

    # Pick 20% of groups (>=1) to be active and give them equal coefficients
    n_active = max(1, n_groups // 5)
    var_per_group = group_size * (1 + (group_size - 1) * corr)
    coeff = np.sqrt(target_irreducible_mse * snr / (n_active * var_per_group))
    for g in range(n_active):
        beta[g * group_size : (g + 1) * group_size] = coeff

    # Build block-diagonal covariance and draw predictors
    cov_block = np.full((group_size, group_size), corr)
    np.fill_diagonal(cov_block, 1.0)
    sigma = block_diag(*([cov_block] * n_groups))
    X = np.random.multivariate_normal(np.zeros(n_features), sigma, n_obs)

    noise_sd = np.sqrt(target_irreducible_mse)
    y = X @ beta + np.random.normal(0, noise_sd, n_obs)

    return X, y, beta, group_indices

def make_colinear_data(
        n_obs: int, group_size: int, n_groups: int, corr: float, seed: int):
    """
    Generates colinear data
    """
    np.random.seed(seed)
    n_features = n_groups * group_size
    beta = np.zeros(n_features)
    beta[:group_size] = np.random.randint(low = 5, high = 15, size = group_size)
    # Build block-diagonal covariance and draw predictors
    cov_block = np.full((group_size, group_size), corr)
    np.fill_diagonal(cov_block, 1.0)
    sigma = block_diag(*([cov_block] * n_groups))
    X = np.random.multivariate_normal(np.zeros(n_features), sigma, n_obs)
    y = X @ beta + np.random.normal(0, 5, n_obs)
    return X, y, beta


# Uncomment to create output directory if it doesn't exist
import os
os.makedirs("../Figures/", exist_ok = True)

# ----------------------------------------------------------------------------
# Figure 1: Paper ------------------------------------------------------------
# ----------------------------------------------------------------------------

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from asgl import Regressor

# Set random seed for reproducibility
np.random.seed(123)

# Generate example data
n = 100000
p = 20
X = np.random.randn(n, p)
true_beta = np.linspace(-4, 4, num = p)
y = X @ true_beta + np.random.randn(n)

lasso_asgl = Regressor(model = "lm", penalization = "lasso", lambda1 = 1)
lasso_asgl.fit(X, y)
lasso_asgl_coefs = lasso_asgl.coef_

alasso_asgl = Regressor(
    model = "lm", 
    penalization = "alasso", 
    weight_technique = "unpenalized", 
    lambda1 = 1)
alasso_asgl.fit(X, y)
alasso_asgl_coefs = alasso_asgl.coef_

# Create DataFrame for plotting
alasso_df = pd.DataFrame({
    "True beta": true_beta,
    "Lasso asgl": lasso_asgl_coefs,
    "Adaptive lasso asgl": alasso_asgl_coefs})

# Melt the DataFrame for plotting
alasso_df_melted = alasso_df.melt(
    id_vars = "True beta",
    var_name = "Penalization",
    value_name = "Predicted beta")

# Plotting
sns.set_theme(style = "whitegrid", font_scale = 1.5)
plt.figure(figsize = (12, 6))

# Plot Lasso and Adaptive Lasso coefficients
sns.lineplot(
    data = alasso_df_melted,
    x = "True beta",
    y = "Predicted beta",
    hue = "Penalization")

# Plot the true beta line
plt.plot(
    true_beta,
    true_beta,
    linestyle = "dashed",
    color = "gray",
    label = "True beta")

# Customize the plot
plt.xlabel("True beta")
plt.ylabel("Predicted beta")
plt.legend()
plt.savefig(
    "../Figures/thresholding_function.pdf",
    format = "pdf", bbox_inches = "tight")
plt.close()

# ----------------------------------------------------------------------------
# Example 1: Lasso penalization + Scikit-Learn integration -------------------
# ----------------------------------------------------------------------------

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from asgl import Regressor
from sklearn.datasets import make_regression
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import GridSearchCV, train_test_split

# Generate synthetic regression data
X, y = make_regression(
    n_samples = 1000, n_features = 20,
    n_informative = 5, bias = 10, noise = 5, random_state = 42)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size = 250, random_state = 42)

# Create a Regressor object for linear regression with Lasso penalization
model = Regressor(model = "lm", penalization = "lasso")

# Define a parameter grid for hyper-parameter tuning
lambda1 =  [1e-4, 1e-3, 1e-2, 1e-1, 1]
param_grid = {"lambda1": lambda1}

# Cross-validation process
gscv = GridSearchCV(
    model, param_grid, scoring = "neg_mean_squared_error", n_jobs = -1, cv = 3)
gscv.fit(X_train, y_train)

sns.set_theme(style = "whitegrid", font_scale = 1.5)
plt.figure(figsize = (12, 6))
plt.plot(lambda1, -1 * gscv.cv_results_["mean_test_score"], marker = "o")
plt.xscale("log")
plt.xlabel("Lambda (log scale)")
plt.ylabel("Mean Squared Error")
plt.title("MSE vs. Lambda1 for Lasso Regression (3-fold CV)")
plt.savefig("../Figures/mse_lasso.pdf", format = "pdf", bbox_inches = "tight")
plt.close()

# Make predictions on the test data
predictions = gscv.predict(X_test)

# Evaluate the model's performance using mean squared error
mse = np.round(mean_squared_error(predictions, y_test), 3)
print(f"Mean Squared Error: {mse}")

# ----------------------------------------------------------------------------
# Example 2: lasso vs adaptive lasso -----------------------------------------
# ----------------------------------------------------------------------------

import time
from datetime import datetime
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from asgl import Regressor
from sklearn.linear_model import Lasso
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import GridSearchCV, train_test_split

n_repetitions = 25
n_groups_arr = np.array([2, 4, 8, 16, 32])

sklearn_lasso_grid = {"alpha": 10 ** np.arange(-1, 1, 0.05)}
asgl_lasso_grid = {"lambda1": 10 ** np.arange(-1, 1, 0.05)}

results = []
idx_seed = 0
for n_groups in n_groups_arr:
    n_features = n_groups * 25
    current_time = datetime.now().strftime("%H:%M:%S")
    print(f"{current_time}. Running simulations for p = {n_features}...")
    for it in range(n_repetitions):
        idx_seed += 1
        X, y, beta_true = make_colinear_data(
            n_obs = 200, group_size = 25, n_groups = n_groups, 
            corr = 0.7, seed = idx_seed)
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size = 100, random_state = it)

        # Lasso using scikit-learn package
        start = time.perf_counter()
        gscv_lasso_sk = GridSearchCV(
            Lasso(), sklearn_lasso_grid, 
            scoring = "neg_mean_squared_error", cv = 3, n_jobs = -1)
        gscv_lasso_sk.fit(X_train, y_train)
        lasso_sk_pred = gscv_lasso_sk.predict(X_test)
        results.append({
            "Iteration": it, 
            "Features": n_features,
            "Model": "sklearn lasso",
            "MSE": mean_squared_error(lasso_sk_pred, y_test),
            "L2 betas": np.linalg.norm(
                beta_true - gscv_lasso_sk.best_estimator_.coef_),
            "Exec time": time.perf_counter() - start})

        # Adaptive lasso using asgl package
        start = time.perf_counter()
        gscv_alasso_asgl = GridSearchCV(
            Regressor(penalization = "alasso"), asgl_lasso_grid, 
            scoring = "neg_mean_squared_error", cv = 3, n_jobs = -1)
        gscv_alasso_asgl.fit(X_train, y_train)
        alasso_asgl_pred = gscv_alasso_asgl.predict(X_test)
        results.append({
            "Iteration": it, 
            "Features": n_features,
            "Model": "asgl adaptive lasso",
            "MSE": mean_squared_error(alasso_asgl_pred, y_test),
            "L2 betas": np.linalg.norm(
                beta_true - gscv_alasso_asgl.best_estimator_.coef_),
            "Exec time": time.perf_counter() - start})

# Convert to dataframe
df_ex2 = pd.DataFrame(
    results, 
    columns = ["Iteration", "Features", "Model", "MSE", "L2 betas", "Exec time"])

# Aggregate across iterations 
mean_df = (df_ex2.groupby(["Features", "Model"], as_index = False)
             .agg({"MSE":"mean", "L2 betas":"mean", "Exec time":"mean"})
             .rename(columns = {"L2 betas":"L2", "Exec time":"Time"}))

# Build the 1×3 panel
import matplotlib.pyplot as plt
fig, axes = plt.subplots(1, 3, figsize = (15, 4), sharex = True)
metrics = ["MSE", "L2", "Time"]

# Plot metrics
for ax, m in zip(axes, metrics):
    for model, d in mean_df.groupby("Model"):
        ax.plot(d["Features"], d[m], marker = "o", label = model)
    ax.set_title(m)
    ax.set_xlabel("Features")
    ax.set_ylabel(f"Mean {m}")

handles, labels = axes[0].get_legend_handles_labels()
fig.legend(handles, labels, loc = "lower center", ncol = len(labels),
    bbox_to_anchor = (0.5, -0.05))

plt.tight_layout(rect = [0, 0.05, 1, 1])
plt.savefig(
    "../Figures/lasso_vs_adaptive_lasso.pdf", 
    format = "pdf", bbox_inches = "tight")
plt.close()

# ----------------------------------------------------------------------------
# Example 3: lasso vs adaptive sparse group lasso in quantile regression  ----
# ----------------------------------------------------------------------------

import time
from datetime import datetime
import numpy as np
import pandas as pd
from asgl import Regressor
from sklearn.linear_model import QuantileRegressor
from sklearn.metrics import median_absolute_error
from sklearn.model_selection import GridSearchCV, train_test_split

n_repetitions = 25
sklearn_lasso_grid = {"alpha": 10 ** np.arange(-2, 1, 0.1)}
asgl_asgl_grid = {
    "lambda1": 10 ** np.arange(-2, 1, 0.1), 
    "alpha": np.arange(0, 1.01, 0.25)}

results = []
for it in range(n_repetitions):
    current_time = datetime.now().strftime("%H:%M:%S")
    print(f"{current_time}. Running iteration: {it+1}/{n_repetitions}")

    X, y, beta_true, group_index = make_group_data(
        n_obs = 200, n_groups = 15, group_size = 15, corr = 0.5, seed = it)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size = 100, random_state = it)

    # Lasso using scikit-learn
    start = time.perf_counter()
    gscv_lasso_sk = GridSearchCV(
        QuantileRegressor(), sklearn_lasso_grid, 
        scoring = "neg_median_absolute_error", cv = 3, n_jobs = -1)
    gscv_lasso_sk.fit(X_train, y_train)
    lasso_sk_pred = gscv_lasso_sk.predict(X_test)
    results.append({
        "Iteration": it, 
        "Model": "scikit-learn lasso",
        "MAE": median_absolute_error(lasso_sk_pred, y_test),
        "L2 betas": np.linalg.norm(
            beta_true - gscv_lasso_sk.best_estimator_.coef_),
        "Exec time": time.perf_counter() - start})

    # Adaptive sparse group lasso using asgl package
    start = time.perf_counter()
    gscv_asgl_asgl = GridSearchCV(
        Regressor(model = "qr", penalization = "asgl", solver = "CLARABEL"), 
        asgl_asgl_grid, scoring = "neg_median_absolute_error", cv = 3, n_jobs = -1)
    gscv_asgl_asgl.fit(X = X_train, y = y_train, group_index = group_index)
    asgl_asgl_pred = gscv_asgl_asgl.predict(X_test)
    results.append({
        "Iteration": it, 
        "Model": "asgl adaptive sgl",
        "MAE": median_absolute_error(asgl_asgl_pred, y_test),
        "L2 betas": np.linalg.norm(
            beta_true - gscv_asgl_asgl.best_estimator_.coef_),
        "Exec time": time.perf_counter() - start})

df_ex3 = pd.DataFrame(
    results, 
    columns = ["Iteration", "Model", "MAE", "L2 betas", "Exec time"])
summary_ex3 = (df_ex3.groupby("Model")
    [["MAE", "L2 betas", "Exec time"]].agg(["mean", "std"]))
pd.set_option("display.max_columns", None)
pd.set_option("display.width", None)
print(summary_ex3)

# ----------------------------------------------------------------------------
# Example 4: Adaptive sparse group lasso linear regression with custom weight-
# ----------------------------------------------------------------------------

import numpy as np
from asgl import Regressor
from sklearn.datasets import make_regression

np.random.seed(1234)

# Generate synthetic regression data
X, y = make_regression(
    n_samples = 200, n_features = 200,
    n_informative = 25, bias = 10, noise = 5, random_state = 42)

# Define the group structure
group_index = np.random.randint(1, 5, size = 200)

# Generate custom weights
custom_individual_weights = np.random.rand(X.shape[1])
custom_group_weights = np.random.rand(len(np.unique(group_index)))

# Create a Regressor object with custom weights
model = Regressor(
    penalization = "asgl",
    individual_weights = custom_individual_weights,
    group_weights = custom_group_weights, solver = "CLARABEL") 

# Fit the model
model.fit(X, y, group_index = group_index)

# -------------------------------------------------------------------------------
# Example 5: Execution time comparison scikit-learn's lasso vs asgl lasso -------
# -------------------------------------------------------------------------------

import time
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from asgl import Regressor
from sklearn.datasets import make_regression
from sklearn.linear_model import Lasso

# Define parameter grids
n_samples_fixed = 200
n_features_grid = [200, 500, 1000, 5000]

n_features_fixed = 200
n_samples_grid = [200, 500, 1000, 5000]

n_repetitions = 20

results_list = []
# Scenario 1: Fixed n_samples, varying n_features
print(f"Running Scenario 1: Fixed n_samples = {n_samples_fixed}, varying n_features")
for n_features_val in n_features_grid:
    print(f"  n_features = {n_features_val}")
    for i in range(1, n_repetitions + 1):
        X, y, beta_true = make_regression(
            n_samples = n_samples_fixed, n_features = n_features_val, 
            n_informative = 25, bias = 10, noise = 5, 
            random_state = i, coef = True)

        # Train scikit-learn Lasso
        start_time = time.perf_counter()
        lasso_sk = Lasso(alpha = 0.05, max_iter = 10000) # Increased max_iter for convergence with more features
        lasso_sk.fit(X, y)
        lasso_sk_exec = time.perf_counter() - start_time
        # Train asgl lasso
        start_time = time.perf_counter()
        lasso_asgl = Regressor(
            lambda1 = 0.05, penalization = "lasso", solver = "CLARABEL")
        lasso_asgl.fit(X, y)
        lasso_asgl_exec = time.perf_counter() - start_time
        results_list.append({
            "sample_size": n_samples_fixed,
            "n_features": n_features_val,
            "repetition": i,
            "scikit-learn": lasso_sk_exec,
            "asgl": lasso_asgl_exec})

sample_df = pd.DataFrame(results_list)

results_list = []
# Scenario 2: Fixed n_features, varying n_samples
print(f"\nRunning Scenario 2: Fixed n_features = {n_features_fixed}, varying n_samples")

for n_samples_val in n_samples_grid:
    print(f"  n_samples = {n_samples_val}")
    for i in range(1, n_repetitions + 1):
        X, y, beta_true = make_regression(
            n_samples = n_samples_val, n_features = n_features_fixed, 
            n_informative = 25, bias = 10, noise = 5, 
            random_state = i, coef = True)

        # Train scikit-learn Lasso
        start_time = time.perf_counter()
        lasso_sk = Lasso(alpha = 0.05, max_iter = 10000) # Increased max_iter for convergence with more features
        lasso_sk.fit(X, y)
        lasso_sk_exec = time.perf_counter() - start_time

        # Train asgl lasso
        start_time = time.perf_counter()
        lasso_asgl = Regressor(
            lambda1 = 0.05, penalization = "lasso", solver = "CLARABEL")
        lasso_asgl.fit(X, y)
        lasso_asgl_exec = time.perf_counter() - start_time

        results_list.append({
            "sample_size": n_samples_val,
            "n_features": n_features_fixed,
            "repetition": i,
            "scikit-learn": lasso_sk_exec,
            "asgl": lasso_asgl_exec
        })

features_df = pd.DataFrame(results_list)

# Plot
sns.set_theme(style = "whitegrid", palette = "colorblind", font_scale = 1.5)
fig, axes = plt.subplots(1, 2, figsize = (12, 5), sharey = True)

# LEFT: fix sample size, vary n_features
left_long = sample_df.melt(
    id_vars = "n_features",
    value_vars = ["scikit-learn", "asgl"],
    var_name = "Method", value_name = "Time")

sns.violinplot(
    data = left_long, x = "n_features", y = "Time", hue = "Method", ax = axes[0],
    palette = "colorblind", fill = True, linecolor = "auto", linewidth = 0.0001,
    inner = None
)

axes[0].set_title(f"Execution Time (Fixed N samples = {n_samples_fixed})")
axes[0].set_ylabel("Execution Time (seconds)")
axes[0].set_xlabel("Number of Features")
axes[0].set_yscale("log")

# RIGHT: fix n_features, vary sample_size
right_long = features_df.melt(
    id_vars = "sample_size",
    value_vars = ["scikit-learn", "asgl"],
    var_name = "Method", value_name = "Time"
)

sns.violinplot(
    data = right_long, x = "sample_size", y = "Time", hue = "Method", ax = axes[1], 
    palette = "colorblind", fill = True, linecolor = "auto", linewidth = 0.0001,
    inner = None
)

axes[1].set_title(f"Execution Time (Fixed N features = {n_features_fixed})")
axes[1].set_ylabel("Execution Time (seconds)")
axes[1].set_xlabel("Number of samples")
axes[1].set_yscale("log")

# Shared legend outside the plotting area
for ax in axes:
    ax.get_legend().remove()

handles, labels = axes[1].get_legend_handles_labels()
fig.legend(handles, labels, loc = "center left",
    bbox_to_anchor = (1.01, 0.5), borderaxespad = 0)

fig.tight_layout()
plt.savefig(
    "../Figures/asgl_vs_sklearn_execution_times.pdf", 
    bbox_inches = "tight")
plt.close()