import numpy as np
import pandas as pd
from numpy.random import default_rng
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import time


# ============================================================
# 1. Data generator: small T, large p, sparse true beta
# ============================================================

def generate_data(
    T=100,
    p=200,
    s=4,
    rho=0.5,
    sigma=1.0,
    seed=None,
    test_size=500
):
    """
    Generate train/test data for a sparse linear model:

        y = X beta0 + eps

    with T observations and p regressors.

    Parameters
    ----------
    T : int
        Training sample size.
    p : int
        Number of candidate regressors.
    s : int
        Number of truly active regressors.
    rho : float
        Correlation parameter for Toeplitz covariance:
        Sigma_ij = rho^{|i-j|}.
    sigma : float
        Error standard deviation.
    seed : int or None
        Random seed.
    test_size : int
        Size of independent test sample.

    Returns
    -------
    X_train, y_train, X_test, y_test, beta0
    """

    rng = default_rng(seed)

    # Toeplitz covariance matrix
    idx = np.arange(p)
    Sigma = rho ** np.abs(np.subtract.outer(idx, idx))

    # Training regressors
    X_train = rng.multivariate_normal(
        mean=np.zeros(p),
        cov=Sigma,
        size=T
    )

    # Test regressors
    X_test = rng.multivariate_normal(
        mean=np.zeros(p),
        cov=Sigma,
        size=test_size
    )

    # Sparse true beta
    beta0 = np.zeros(p)
    base_vals = np.array([2.0, 1.5, -1.0, 0.8, -0.6, 0.5, -0.4, 0.3])

    if s <= len(base_vals):
        beta0[:s] = base_vals[:s]
    else:
        beta0[:len(base_vals)] = base_vals
        beta0[len(base_vals):s] = rng.normal(0.0, 0.3, size=s-len(base_vals))

    # Errors
    eps_train = rng.normal(0.0, sigma, size=T)
    eps_test = rng.normal(0.0, sigma, size=test_size)

    y_train = X_train @ beta0 + eps_train
    y_test = X_test @ beta0 + eps_test

    return X_train, y_train, X_test, y_test, beta0


# ============================================================
# 2. Criterion functions
# ============================================================

def rss_score(y, yhat):
    """
    Residual sum of squares.
    """
    return np.sum((y - yhat) ** 2)


def aic_score(y, yhat, k):
    """
    Gaussian linear-model AIC up to constants:

        AIC = n log(RSS/n) + 2k
    """
    n = len(y)
    rss = rss_score(y, yhat)
    rss = max(rss, 1e-12)
    return n * np.log(rss / n) + 2 * k


def bic_score(y, yhat, k):
    """
    Gaussian linear-model BIC up to constants:

        BIC = n log(RSS/n) + k log(n)
    """
    n = len(y)
    rss = rss_score(y, yhat)
    rss = max(rss, 1e-12)
    return n * np.log(rss / n) + k * np.log(n)


def criterion_value(y, yhat, k, criterion="rss"):
    """
    Compute criterion value.

    Lower is better.
    """
    if criterion == "rss":
        return rss_score(y, yhat)
    elif criterion == "aic":
        return aic_score(y, yhat, k)
    elif criterion == "bic":
        return bic_score(y, yhat, k)
    else:
        raise ValueError("criterion must be one of: 'rss', 'aic', 'bic'")


# ============================================================
# 3. Fit OLS on a selected subset
# ============================================================

def fit_subset_ols(X, y, selected):
    """
    Fit OLS using only selected regressors.

    Returns
    -------
    model : sklearn LinearRegression or None
    yhat : fitted values
    beta_hat_subset : coefficients on selected regressors
    """

    n = len(y)

    if len(selected) == 0:
        ybar = np.mean(y)
        yhat = np.full(n, ybar)
        return None, yhat, np.array([])

    model = LinearRegression(fit_intercept=True)
    model.fit(X[:, selected], y)
    yhat = model.predict(X[:, selected])

    return model, yhat, model.coef_


# ============================================================
# 4. Forward selection with threshold stopping
# ============================================================

def forward_selection(
    X,
    y,
    tau=1e-3,
    criterion="rss",
    max_features=None,
    min_relative_improvement=False,
    verbose=False
):
    """
    Forward selection algorithm.

    Start from empty model. At each step, add the variable that produces
    the largest criterion improvement. Stop when improvement < tau.

    Parameters
    ----------
    X : ndarray, shape (T, p)
        Regressor matrix.
    y : ndarray, shape (T,)
        Outcome vector.
    tau : float
        Stopping threshold.
        If min_relative_improvement=False:
            stop when absolute improvement < tau.
        If min_relative_improvement=True:
            stop when relative improvement < tau.
    criterion : {'rss', 'aic', 'bic'}
        Criterion to minimize.
    max_features : int or None
        Maximum selected regressors. If None, uses min(T-1, p).
    min_relative_improvement : bool
        Whether threshold is relative to previous criterion.
    verbose : bool
        Print progress.

    Returns
    -------
    selected : list[int]
        Selected regressors.
    history : pd.DataFrame
        Selection path.
    """

    T, p = X.shape

    if max_features is None:
        # Do not allow too many regressors relative to sample size.
        # Since LinearRegression has an intercept, keep selected <= T-1.
        max_features = min(p, max(1, T - 1))

    remaining = list(range(p))
    selected = []

    # Empty model
    _, yhat_empty, _ = fit_subset_ols(X, y, selected)
    current_value = criterion_value(y, yhat_empty, k=1, criterion=criterion)

    history_rows = []

    step = 0

    while len(remaining) > 0 and len(selected) < max_features:

        best_var = None
        best_value = current_value
        best_yhat = None

        for j in remaining:
            trial = selected + [j]
            _, yhat_trial, _ = fit_subset_ols(X, y, trial)

            # k includes intercept + selected regressors
            k_trial = 1 + len(trial)
            val = criterion_value(y, yhat_trial, k=k_trial, criterion=criterion)

            if val < best_value:
                best_value = val
                best_var = j
                best_yhat = yhat_trial

        if best_var is None:
            if verbose:
                print("No variable improves the criterion.")
            break

        absolute_improvement = current_value - best_value

        if min_relative_improvement:
            denom = max(abs(current_value), 1e-12)
            improvement_for_stop = absolute_improvement / denom
        else:
            improvement_for_stop = absolute_improvement

        if verbose:
            print(
                f"Step {step+1}: add {best_var}, "
                f"criterion {current_value:.4f} -> {best_value:.4f}, "
                f"improvement={improvement_for_stop:.6f}"
            )

        if improvement_for_stop < tau:
            if verbose:
                print("Stopping: improvement below threshold.")
            break

        step += 1
        selected.append(best_var)
        remaining.remove(best_var)

        history_rows.append({
            "step": step,
            "added_var": best_var,
            "criterion_before": current_value,
            "criterion_after": best_value,
            "absolute_improvement": absolute_improvement,
            "stopping_improvement": improvement_for_stop,
            "selected_size": len(selected)
        })

        current_value = best_value

    history = pd.DataFrame(history_rows)
    return selected, history


# ============================================================
# 5. Evaluation metrics
# ============================================================

def evaluate_selection(
    X_train,
    y_train,
    X_test,
    y_test,
    beta0,
    selected,
    true_support
):
    """
    Evaluate selected model.

    Returns dictionary with:
    - selected size
    - true positives
    - false positives
    - false negatives
    - exact support recovery
    - beta L2 error
    - in-sample MSE
    - out-of-sample MSE
    """

    p = X_train.shape[1]

    model, yhat_train, beta_subset = fit_subset_ols(X_train, y_train, selected)

    beta_hat = np.zeros(p)

    if len(selected) > 0:
        beta_hat[selected] = beta_subset
        yhat_test = model.predict(X_test[:, selected])
    else:
        yhat_test = np.full(len(y_test), np.mean(y_train))

    selected_set = set(selected)
    true_set = set(true_support)

    tp = len(selected_set.intersection(true_set))
    fp = len(selected_set.difference(true_set))
    fn = len(true_set.difference(selected_set))

    exact_recovery = int(selected_set == true_set)

    return {
        "selected_size": len(selected),
        "tp": tp,
        "fp": fp,
        "fn": fn,
        "exact_recovery": exact_recovery,
        "beta_l2_error": np.linalg.norm(beta_hat - beta0),
        "train_mse": mean_squared_error(y_train, yhat_train),
        "test_mse": mean_squared_error(y_test, yhat_test)
    }


# ============================================================
# 6. Monte Carlo experiment
# ============================================================

def monte_carlo_forward_selection(
    R=200,
    T=100,
    p=200,
    s=4,
    rho=0.5,
    sigma=1.0,
    tau=0.01,
    criterion="rss",
    max_features=None,
    min_relative_improvement=False,
    test_size=500,
    seed=2026,
    verbose=False
):
    """
    Monte Carlo driver for forward selection.

    Parameters
    ----------
    R : int
        Number of replications.
    T : int
        Training sample size.
    p : int
        Number of regressors.
    s : int
        True sparsity.
    rho : float
        Covariance correlation parameter.
    sigma : float
        Error standard deviation.
    tau : float
        Stopping threshold.
    criterion : {'rss', 'aic', 'bic'}
        Selection criterion.
    max_features : int or None
        Maximum number of selected variables.
    min_relative_improvement : bool
        Whether to use relative improvement stopping.
    test_size : int
        Test sample size.
    seed : int
        Master seed.
    verbose : bool
        Print progress.

    Returns
    -------
    results : pd.DataFrame
        Monte Carlo results.
    histories : list[pd.DataFrame]
        Selection histories for each replication.
    """

    rng = default_rng(seed)
    rows = []
    histories = []

    start_time = time.time()

    true_support = list(range(s))

    for r in range(R):

        rep_seed = rng.integers(0, 10**9)

        X_train, y_train, X_test, y_test, beta0 = generate_data(
            T=T,
            p=p,
            s=s,
            rho=rho,
            sigma=sigma,
            seed=rep_seed,
            test_size=test_size
        )

        selected, history = forward_selection(
            X_train,
            y_train,
            tau=tau,
            criterion=criterion,
            max_features=max_features,
            min_relative_improvement=min_relative_improvement,
            verbose=False
        )

        metrics = evaluate_selection(
            X_train,
            y_train,
            X_test,
            y_test,
            beta0,
            selected,
            true_support
        )

        metrics.update({
            "rep": r,
            "T": T,
            "p": p,
            "s": s,
            "rho": rho,
            "sigma": sigma,
            "tau": tau,
            "criterion": criterion,
            "relative_stop": min_relative_improvement,
            "selected_vars": selected
        })

        rows.append(metrics)
        histories.append(history)

        if verbose and (r + 1) % 25 == 0:
            print(f"Completed {r+1}/{R} replications.")

    elapsed = time.time() - start_time
    results = pd.DataFrame(rows)
    results.attrs["elapsed_seconds"] = elapsed

    return results, histories


# ============================================================
# 7. Summary utilities
# ============================================================

def summarize_results(results):
    """
    Summarize Monte Carlo results.
    """

    summary = {
        "R": len(results),
        "mean_selected_size": results["selected_size"].mean(),
        "median_selected_size": results["selected_size"].median(),
        "mean_tp": results["tp"].mean(),
        "mean_fp": results["fp"].mean(),
        "mean_fn": results["fn"].mean(),
        "exact_recovery_rate": results["exact_recovery"].mean(),
        "mean_beta_l2_error": results["beta_l2_error"].mean(),
        "mean_train_mse": results["train_mse"].mean(),
        "mean_test_mse": results["test_mse"].mean(),
        "elapsed_seconds": results.attrs.get("elapsed_seconds", np.nan)
    }

    return pd.Series(summary)


def print_summary(results):
    """
    Print readable summary.
    """

    summary = summarize_results(results)

    print("\n==============================")
    print("Forward Selection MC Summary")
    print("==============================")
    print(summary.to_string())
    print("==============================\n")


# ============================================================
# 8. Threshold comparison experiment
# ============================================================

def compare_thresholds(
    taus=(0.001, 0.01, 0.05, 0.1),
    R=100,
    T=100,
    p=200,
    s=4,
    rho=0.5,
    sigma=1.0,
    criterion="rss",
    min_relative_improvement=False,
    seed=2026
):
    """
    Run the Monte Carlo for several stopping thresholds.
    """

    summaries = []

    for tau in taus:
        results, _ = monte_carlo_forward_selection(
            R=R,
            T=T,
            p=p,
            s=s,
            rho=rho,
            sigma=sigma,
            tau=tau,
            criterion=criterion,
            min_relative_improvement=min_relative_improvement,
            seed=seed + int(1e6 * tau)
        )

        summary = summarize_results(results)
        summary["tau"] = tau
        summary["criterion"] = criterion
        summaries.append(summary)

    return pd.DataFrame(summaries)


# ============================================================
# 9. Criterion comparison experiment
# ============================================================

def compare_criteria(
    criteria=("rss", "aic", "bic"),
    R=100,
    T=100,
    p=200,
    s=4,
    rho=0.5,
    sigma=1.0,
    tau=0.01,
    min_relative_improvement=False,
    seed=2026
):
    """
    Compare RSS, AIC, and BIC stopping rules.
    """

    summaries = []

    for c in criteria:
        results, _ = monte_carlo_forward_selection(
            R=R,
            T=T,
            p=p,
            s=s,
            rho=rho,
            sigma=sigma,
            tau=tau,
            criterion=c,
            min_relative_improvement=min_relative_improvement,
            seed=seed + hash(c) % 10_000
        )

        summary = summarize_results(results)
        summary["criterion"] = c
        summary["tau"] = tau
        summaries.append(summary)

    return pd.DataFrame(summaries)


# ============================================================
# 10. Main script
# ============================================================

if __name__ == "__main__":

    # --------------------------------------------------------
    # Baseline experiment
    # --------------------------------------------------------

    results, histories = monte_carlo_forward_selection(
        R=200,
        T=100,
        p=200,
        s=4,
        rho=0.5,
        sigma=1.0,
        tau=0.01,
        criterion="rss",
        max_features=20,
        min_relative_improvement=True,
        test_size=500,
        seed=2026,
        verbose=True
    )

    print_summary(results)

    print("First few rows:")
    print(results.head())

    # --------------------------------------------------------
    # Threshold comparison
    # --------------------------------------------------------

    threshold_table = compare_thresholds(
        taus=(0.001, 0.005, 0.01, 0.05),
        R=100,
        T=100,
        p=200,
        s=4,
        rho=0.5,
        sigma=1.0,
        criterion="rss",
        min_relative_improvement=True,
        seed=2026
    )

    print("\nThreshold comparison:")
    print(threshold_table[
        [
            "tau",
            "mean_selected_size",
            "mean_tp",
            "mean_fp",
            "exact_recovery_rate",
            "mean_beta_l2_error",
            "mean_test_mse"
        ]
    ])

    # --------------------------------------------------------
    # Criterion comparison
    # --------------------------------------------------------

    criterion_table = compare_criteria(
        criteria=("rss", "aic", "bic"),
        R=100,
        T=100,
        p=200,
        s=4,
        rho=0.5,
        sigma=1.0,
        tau=0.01,
        min_relative_improvement=True,
        seed=2026
    )

    print("\nCriterion comparison:")
    print(criterion_table[
        [
            "criterion",
            "mean_selected_size",
            "mean_tp",
            "mean_fp",
            "exact_recovery_rate",
            "mean_beta_l2_error",
            "mean_test_mse"
        ]
    ])