import shutil
from importlib.resources import as_file, files
from pathlib import Path

from scipy.stats import norm
from sklearn.linear_model import LinearRegression

# REGENIE input files preparation helper functions
def rint_normalization(ser):
    ranks = ser.rank()
    normalized = norm.ppf((ranks - 0.5)/ranks.notna().sum())
    return normalized

def save_files_in_regenie_fmt(
    df, save_file, sample_id_col="person_id",
    phenos=['bmi_rint', 'whradjbmi_rint'],
    covariates=[
        'genetic_sex', 'age', 'age_2', 'age_sex', 'ancestry_pred_other',
        'genetic_pca1', 'genetic_pca2', 'genetic_pca3', 'genetic_pca4', 'genetic_pca5', 
        'genetic_pca6', 'genetic_pca7', 'genetic_pca8', 'genetic_pca9', 'genetic_pca10'
    ]):
    curr_columns = df.columns
    df["FID"] = df[sample_id_col]
    df["IID"] = df[sample_id_col]
    df.loc[:, ["FID", "IID"] + phenos + covariates].fillna("NA").to_csv(save_file, index=False, sep="\t")
    return

def get_residuals(X, y):
    model = LinearRegression()
    model.fit(X, y)
    predictions = model.predict(X)
    residuals = y - predictions
    return residuals

# REGENIE association test helper functions
def get_regenie_wdl_resource():
    return files("biobanking").joinpath("workflows/regenie.wdl")

def copy_regenie_wdl(destination_dir):
    destination_dir = Path(destination_dir)
    destination_dir.mkdir(parents=True, exist_ok=True)
    with as_file(get_regenie_wdl_resource()) as source_path:
        destination_path = destination_dir / "regenie.wdl"
        shutil.copy2(source_path, destination_path)
    return destination_path

def list_gcs_files(bucket_name, prefix):
    from google.cloud import storage

    storage_client = storage.Client()
    bucket_obj = storage_client.bucket(bucket_name.lstrip("gs://"))
    blobs = bucket_obj.list_blobs(prefix=prefix)
    file_paths = [f"{bucket_name}/{blob.name}" for blob in blobs]
    return file_paths
