import numpy as np
import pandas as pd
from sklearn.linear_model import LassoCV, ElasticNetCV
from scipy.cluster.hierarchy import linkage, fcluster
from scipy.spatial.distance import squareform
from sklearn.preprocessing import StandardScaler
from statsmodels.stats.outliers_influence import variance_inflation_factor

def correlation_clustering(X, threshold=0.85):
    corr = X.corr().abs()
    dist = 1 - corr
    Z = linkage(squareform(dist.values), method="average")
    clusters = fcluster(Z, t=1 - threshold, criterion="distance")

    keep = []
    dfc = pd.DataFrame({"col": X.columns, "cluster": clusters})

    for cl in dfc.cluster.unique():
        members = dfc[dfc.cluster == cl].col.tolist()
        if len(members) == 1:
            keep.append(members[0])
        else:
            stds = X[members].std()
            keep.append(stds.idxmax())

    dropped = [c for c in X.columns if c not in keep]
    return X[keep], dropped

def compute_vif(X):
    X = X.copy().assign(const=1)
    vif = pd.Series(
        [variance_inflation_factor(X.values, i) for i in range(X.shape[1] - 1)],
        index=X.columns[:-1]
    )
    return vif.sort_values(ascending=False)

def vif_filter(X, threshold=10):
    removed = []
    Xw = X.copy()
    while True:
        vif = compute_vif(Xw)
        if vif.iloc[0] <= threshold:
            break
        col = vif.index[0]
        removed.append(col)
        Xw = Xw.drop(columns=[col])
        if Xw.shape[1] <= 2:
            break
    return Xw, removed

def lasso_elastic_selection(X, y):
    scaler = StandardScaler()
    Xs = scaler.fit_transform(X)

    lasso = LassoCV(cv=5, n_jobs=-1).fit(Xs, y)
    enet = ElasticNetCV(l1_ratio=[0.1, 0.5, 0.9], cv=5, n_jobs=-1).fit(Xs, y)

    selected = list(set(
        X.columns[lasso.coef_ != 0].tolist() +
        X.columns[enet.coef_ != 0].tolist()
    ))
    return selected



