# Copyright 2019 The Glow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools
import math
import re
import warnings
from typing import Any, Dict, Iterable, List, Tuple

import numpy as np
import pandas as pd
import pyspark.sql.functions as f
import scipy.optimize
from nptyping import Float, Int, NDArray
from pyspark.sql import DataFrame, Row
from pyspark.sql.window import Window
from typeguard import typechecked


@typechecked
def sort_in_place(pdf: pd.DataFrame, columns: List[str]) -> None:
    """
    A faster alternative to DataFrame.sort_values. Note that this function is less sophisticated
    than sort_values and does not allow for control over sort direction or null handling.

    Adapted from https://github.com/pandas-dev/pandas/issues/15389.

    Args:
        pdf : The pandas DataFrame to sort
        columns : Columns to sort by
    """
    order = np.lexsort([pdf[col].array for col in reversed(columns)])
    for col in list(pdf.columns):
        pdf[col].array[:] = pdf[col].array[order]


@typechecked
def parse_header_block_sample_block_label(key: Tuple,
                                          key_pattern: List[str]) -> Tuple[str, str, str]:
    """
    Extracts header_block, sample_block, and label from the key corresponding to a group from the groupBy(key_pattern)
    clause.  Key pattern can be one of:
        (header_block, sample_block),
        (header_block, sample_block, label),
        (sample_block, label)
    depending on the context.  In each case, a tuple with 3 members is returned, with the missing member filled in by
    'all' where necessary

    Args:
        key : key for the group
        key_pattern : one of the aforementioned key patterns

    Returns:
        tuple of (header_block, sample_block, label), where header_block or label may be filled with 'all'
        depending on context.
    """
    if 'label' in key_pattern:
        label = key[-1]
    else:
        label = 'all'
    if 'header_block' in key_pattern:
        header_block = key[0]
        sample_block = key[1]
    else:
        header_block = 'all'
        sample_block = key[0]

    return header_block, sample_block, label


@typechecked
def parse_header_block_sample_block_label_alpha_name(
        key: Tuple, key_pattern: List[str]) -> Tuple[str, str, str, str]:
    """
    Extracts header_block, sample_block, label, and alpha_name from the key corresponding to a group from the
    groupBy(key_pattern) clause.  Key pattern can be one of:
        (header_block, sample_block, label, alpha_name),
        (sample_block, label, alpha_name)
    depending on the context.  In each case, a tuple with 4 members is returned, with header_block filled with 'all' if
    it is absent from they key

    Args:
        key : key for the group
        key_pattern : one of the aforementioned key patterns

    Returns:
        tuple of (header_block, sample_block, label, alpha_name), where header_block or label may be filled with 'all'
        depending on context.
    """
    if 'header_block' in key_pattern:
        return key
    else:
        header_block = 'all'
        return header_block, key[0], key[1], key[2]


# @typechecked -- typeguard does not support numpy array
def assemble_block(n_rows: Int, n_cols: Int, pdf: pd.DataFrame, cov_matrix: NDArray[(Any, Any),
                                                                                    Float],
                   row_mask: NDArray[Any]) -> NDArray[Float]:
    """
    Creates a dense n_rows by n_cols matrix from the array of either sparse or dense vectors in the Pandas DataFrame
    corresponding to a group.  This matrix represents a block.

    Args:
         n_rows : The number of rows in the resulting matrix
         n_cols : The number of columns in the resulting matrix
         pdf : Pandas DataFrame corresponding to a group
         cov_matrix: 2D numpy array representing covariate columns that should be prepended to matrix X from the block.  Can be
            empty if covariates are not being applied.
        row_mask:  1D numpy array of size n_rows containing booleans used to mask rows from the block X before
            return.

    Returns:
        Dense n_rows - n_masked by n_columns matrix where the columns have been 0-centered and standard scaled.
    """
    mu = pdf['mu'].to_numpy()
    sig = pdf['sig'].to_numpy()

    if 0 in sig:
        raise ValueError(f'Standard deviation cannot be 0.')

    if row_mask.size == 0:
        row_mask = np.full(n_rows, True)

    if 'indices' not in pdf.columns:
        X_raw = np.column_stack(pdf['values'].array)
    else:
        X_raw = np.zeros([n_rows, n_cols])
        for column, row in enumerate(pdf[['indices', 'values']].itertuples()):
            X_raw[row.indices, column] = row.values

    X = ((X_raw - mu) / sig)

    if cov_matrix.any():
        return np.column_stack((cov_matrix, X))[row_mask, :]
    else:
        return X[row_mask, :]


# @typechecked -- typeguard does not support numpy array
def constrained_logistic_fit(X: NDArray[Float], y: NDArray[Float], alpha_arr: NDArray[Float],
                             guess: NDArray[Float], n_cov: Int) -> scipy.optimize.OptimizeResult:
    """
    Fits a logistic regression model using design matrix X and binary label y, with ridge penalization specified
    by the penalty parameters in alpha_arr (with one entry per column in X).  The optimization procedure starts from
    the initial state specified in the guess array.  The first n_cov columns in X are assumed to be fixed effects and
    the corresponding parameters are frozen with their initial values in guess.

    Args:
         X: [n_row, n_col] design matrix
         y: [n_row] binary label vector
         alpha_arr : array of n_col ridge penalty values
         guess : array of n_coll initial parameter guesses
         n_cov : number of columns in X (starting from the left) to consider fixed, and freeze at the initial guess
            during optimization.

    Returns:
        scipy optimize result containing optimized paramters from the fit.
    """
    if guess.size == 0:
        guess = np.zeros(X.shape[1])

    def objective(beta):
        z = X @ beta
        p = 1 / (1 + np.exp(-z))
        eps = 1E-15
        ll = (y * np.log(p + eps) + (1 - y) * np.log(1 - p + eps)).sum()
        return -(ll - np.dot(alpha_arr * beta, beta) / 2) / y.size

    def gradient(beta):
        z = X @ beta
        p = 1 / (1 + np.exp(-z))
        grad = -(np.dot((y - p), X) - beta * alpha_arr) / y.size
        if n_cov > 0:
            grad[:n_cov] = 0
        return grad

    return scipy.optimize.minimize(objective, guess, jac=gradient, method='L-BFGS-B')


# @typechecked -- typeguard does not support numpy array
def get_irls_pieces(X: NDArray[Float], y: NDArray[Float], alpha_value: Float,
                    beta_cov: NDArray[Float]) -> (NDArray[Float], NDArray[Float], NDArray[Float]):
    """
    Fits a logistic regression model logit(p(y|X)) ~ X*beta with L2 shrinkage C = 1/alpha, using scikit-learn.  Returns
    the vector beta, the matrix transpose(X)*diag(p(1-p))*X, and the vector transpose(X)*(y - p), which are the pieces
    needed to perform one step of the IRLS algorithm.

    Args:
         X : Matrix of n samples m features.  Matrix X is expected to include a constant intercept = 1 term as the first
          column.
         y : Binary response vector of  length n
         alpha_value : Shrinkage parameter to be used in the fit
         beta_cov : array of fixed covariate parameter estimates to use in the fitting.

    Returns:
        beta : m length array representing coefficients found from the fit, including the intercept term as the first
        element.
        XtGX : m by m matrix representing transpose(X)*diag(p(1-p))*X.
        XtY : m length array representing transpose(X)*(y - p)

    """
    n_cov = beta_cov.size
    # If we have no observations in this block (i.e, y.sum() == 0), then we should not try to fit a model and instead
    # just return a parameterless model based on the population frequency of observations p0
    if y.sum() > 0:
        alpha_arr = np.zeros(X.shape[1])
        alpha_arr[n_cov:] = alpha_value
        guess = np.zeros(X.shape[1])
        guess[:n_cov] = beta_cov
        ridge_fit = constrained_logistic_fit(X, y, alpha_arr, guess, n_cov)
        beta = ridge_fit.x
    else:
        beta = np.zeros(X.shape[1])
        if n_cov > 0:
            beta[:n_cov] = beta_cov

    z = X @ beta
    p = sigmoid(z)
    XtGX = (X.T * (p * (1 - p))) @ X
    XtY = X.T @ (p - y)
    return beta, XtGX, XtY


# @typechecked -- typeguard does not support numpy array
def slice_label_rows(labeldf: pd.DataFrame, label: str, sample_list: List[str],
                     row_mask: NDArray[Any]) -> NDArray[Any]:
    """
    Selects rows from the Pandas DataFrame of labels corresponding to the samples in a particular sample_block.

    Args:
        labeldf : Pandas DataFrame containing the labels
        label : Header for the particular label to slice.  Can be 'all' if all labels are desired.
        sample_list : List of sample ids corresponding to the sample_block to be sliced out.
        row_mask : 1D numpy array of size n_rows containing booleans used to mask samples from the rows sliced from
            labeldf.

    Returns:
        Matrix of [number of samples in sample_block - number of samples masked] x [number of labels to slice]
    """
    if row_mask.size == 0:
        row_mask = np.full(len(sample_list), True)
    if label == 'all':
        return labeldf.loc[sample_list, :].to_numpy()[row_mask, :]
    else:
        return labeldf[label].loc[sample_list].to_numpy().reshape(-1, 1)[row_mask, :]


@typechecked
def evaluate_coefficients(pdf: pd.DataFrame, alpha_values: Iterable[Float],
                          n_cov: int) -> NDArray[Float]:
    """
    Solves the system (XTX + Ia)^-1 * XtY for each of the a values in alphas.  Returns the resulting coefficients.

    Args:
         pdf : Pandas DataFrame for the group
         alpha_values : Array of alpha values (regularization strengths)
         n_cov: Number of covariate columns on the left-most side of matrix X.  These are regularized with a constant
         value of alpha = 1, regardless of the alpha value being used for the rest of the matrix X.

    Returns:
        Matrix of coefficients of size [number of columns in X] x [number of labels * number of alpha values]
    """
    XtX = np.stack(pdf['xtx'].array)
    XtY = np.stack(pdf['xty'].array)
    diags = [
        np.concatenate([np.ones(n_cov), np.ones(XtX.shape[1] - n_cov) * a]) for a in alpha_values
    ]
    return np.column_stack([(np.linalg.inv(XtX + np.diag(d)) @ XtY) for d in diags])


@typechecked
def irls_one_step(pdf: pd.DataFrame, alpha_value: Float, n_cov: Int) -> NDArray[Float]:
    """
    Performs one step of the IRLS algorithm using the components found in pdf and a value of alpha.
    Returns the resulting coefficient values.  If n_cov > 0, then the first n_cov columns from the design matrix X
    are assumed to be fixed effects and the corresponding parameter estimates are not updated during the step.

    Args:
         pdf : Pandas DataFrame for the group.  Contains one set of IRLS equation components (beta, xtgx, xty)
         corresponding to the value of alpha provided
         alpha_value : regularization parameter alpha
         n_cov : number of columns in X (starting from the left) to consider fixed, and freeze at the initial guess
            prior to taking the step.

    Returns:
        vector of coefficients of size [number of columns in X]

    """
    xtgx = np.stack(pdf['xtgx'])[n_cov:, n_cov:]
    xty = pdf['xty'].to_numpy()[n_cov:]
    beta0 = pdf['beta'].to_numpy()
    alpha_arr = np.ones(xtgx.shape[1]) * alpha_value
    dbeta = np.zeros(beta0.size)
    dbeta[n_cov:] = np.linalg.inv(xtgx + np.diag(alpha_arr)) @ (xty + beta0[n_cov:] * alpha_arr)
    return beta0 - dbeta


@typechecked
def cross_alphas_and_labels(alpha_names: Iterable[str], labeldf: pd.DataFrame,
                            label: str) -> List[Tuple[str, str]]:
    """
    Crosses all label and alpha names. The output tuples appear in the same order as the output of
    evaluate_coefficients.

    Args:
        alpha_names : List of string identifiers assigned to the values of alpha
        labeldf : Pandas DataFrame of labels
        label : Label used for this set of coefficients.  Can be 'all' if all labels were used.
    Returns:
        List of tuples of the form (alpha_name, label_name)
    """
    if label == 'all':
        label_names = labeldf.columns
    else:
        label_names = [label]

    return list(itertools.product(alpha_names, label_names))


@typechecked
def new_headers(header_block: str, alpha_names: Iterable[str],
                row_indexer: List[Tuple[str, str]]) -> Tuple[str, List[int], List[str]]:
    """
    Creates new headers for the output of a matrix reduction step.  Generally produces names like
    "block_[header_block_number]_alpha_[alpha_name]_label_[label_name]"

    Args:
        header_block : Identifier for a header_block (e.g., 'chr_1_block_0')
        alpha_names : List of string identifiers for alpha parameters
        row_indexer : A list of tuples provided by the cross_alphas_and_labels function

    Returns:
        new_header_block : A new header_block name, typically the chromosome (e.g. chr_1), but might be 'all' if
        there are no more levels to reduce over.
        sort_keys : Array of sortable integers to specify the ordering of the new matrix headers.
        headers : List of new matrix headers.
    """
    match_chr_block = re.search(r"^chr_(.*)_block_([0-9]+)$", header_block)
    match_chr = re.search(r"^chr_(.*)$", header_block)
    if match_chr_block:
        chr = match_chr_block.group(1)
        inner_index = int(match_chr_block.group(2))
        header_prefix = f'chr_{chr}_block_{inner_index}_'
        new_header_block = f'chr_{chr}'
    elif match_chr:
        chr = match_chr.group(1)
        inner_index = abs(hash(chr)) % (10**8)  # Hash to 8 digits
        header_prefix = f'chr_{chr}_'
        new_header_block = 'all'
    elif header_block == 'all':
        inner_index = 0
        header_prefix = ''
        new_header_block = 'all'
    else:
        raise ValueError(f'Header block {header_block} does not match expected pattern.')

    sort_keys, headers = [], []
    for a, l in row_indexer:
        sort_key = inner_index * len(alpha_names) + int(re.search(r"^alpha_([0-9]+)", a).group(1))
        sort_keys.append(sort_key)
        header = f'{header_prefix}{a}_label_{l}'
        headers.append(header)

    return new_header_block, sort_keys, headers


# @typechecked -- typeguard does not support numpy array
def r_squared(XB: NDArray[Float], Y: NDArray[Float]) -> NDArray[(Any, ), Float]:
    """
    Computes the coefficient of determination (R2) metric between the matrix resulting from X*B and the matrix of labels
    Y.

    Args:
        XB : Matrix representing the result of the multiplication X*B, where X is a matrix of [number of samples] x
        [number of headers] and B is a matrix of [number of headers x number of alphas * number of labels]
        Y : Matrix of labels, with [number of samples x number of labels]

    Returns:
        Array of [number of alphas * number of labels]
    """
    tot = np.power(Y - Y.mean(), 2).sum()
    res = np.power(Y - XB, 2).sum(axis=0)
    return 1 - (res / tot)


# @typechecked -- typeguard does not support numpy array
def sigmoid(z: NDArray[Float]) -> NDArray[Float]:
    """
    Computes the sigmoid function for each element in input z

    Args:
        z: Input values

    Returns:
        Sigmoid response corresponding to each input value
    """
    return 1 / (1 + np.exp(-z))


# @typechecked -- typeguard does not support numpy array
def log_loss(p: NDArray[Float], y: NDArray[Float]) -> NDArray[Float]:
    """
    Computes the log loss of probability values p and observed binary variable y.

    Args:
        p : probability values, assumed to have shape n obervations x m models
        y : observed binary variable, assumed to have shape n observations x 1

    Returns:
        Array of [n * m]
    """
    eps = 1E-15
    return -(y * np.log(p + eps) + (1 - y) * np.log(1 - p + eps)).sum(axis=0) / y.shape[0]


@typechecked
def create_alpha_dict(alphas: List[float]) -> Dict[str, Float]:
    """
    Creates a mapping to attach string identifiers to alpha values.

    Args:
        alphas : Alpha values

    Returns:
        Dict of [alpha names, alpha values]
    """
    if not all(x >= 0 for x in alphas):
        raise Exception('Alpha values must all be non-negative.')
    return {f'alpha_{i}': a for i, a in enumerate(alphas)}


@typechecked
def generate_alphas(blockdf: DataFrame) -> Dict[str, Float]:
    """
    Generates alpha values using a range of heritability values and the number of distinct headers (without labels).

    Args:
        blockdf : Spark DataFrame representing a block matrix
        labeldf: Pandas DataFrame containing target labels

    Returns:
        Dict of [alpha names, alpha values]
    """
    num_label_free_headers = blockdf.select(f.regexp_extract('header', r"^(.+?)($|_label_.*)",
                                                             1)).distinct().count()
    heritability_vals = [0.99, 0.75, 0.50, 0.25, 0.01]
    alphas = [num_label_free_headers / h for h in heritability_vals]
    print(f"Generated alphas: {alphas}")
    return create_alpha_dict(alphas)


@typechecked
def _assert_all_present(df: pd.DataFrame, col_name: Any, df_name: str) -> None:
    """
    Raises an error if a pandas series has missing values.

    Args:
        df : Pandas DataFrame
    """
    if df[col_name].isnull().any():
        raise ValueError(
            f"Missing values are present in the {df_name} dataframe's {col_name} column")


@typechecked
def _is_zero(f: float) -> bool:
    return math.isclose(f, 0, abs_tol=0.01)


@typechecked
def _is_one(f: float) -> bool:
    return math.isclose(f, 1, abs_tol=0.01)


@typechecked
def _num_non_binary_values(s: pd.Series) -> int:
    """
    Returns the number of values in a series that are neither 0, 1, nor missing.
    """
    non_binary_vals = (~s.dropna().isin([0, 1])).sum()
    return int(non_binary_vals)


def _check_binary(df: pd.DataFrame) -> None:
    """
    Warns if any column of a pandas DataFrame is not a binary value equal to 0 or 1.

    Args:
        df : Pandas DataFrame
    """
    for label in df:
        num_non_binary_vals = _num_non_binary_values(df[label])
        if num_non_binary_vals != 0:
            warnings.warn(
                f"Column {label} is not binary. It has {num_non_binary_vals} non-binary value(s).",
                UserWarning)


@typechecked
def _is_binary(df: pd.DataFrame) -> bool:
    """
    Checks whether a Pandas DataFrame is all binary.

    Args:
        df : Pandas DataFrame
    """
    return df.isin([0, 1, None]).all(axis=None).item()


@typechecked
def _fill_na_and_standardize(pdf: pd.DataFrame) -> pd.DataFrame:
    """
    Replaces the NaNs in a Pandas DataFrame with the mean of the column and
    mean centers and scales the columns.

    Args:
        df : Pandas DataFrame
    """
    pdf_mean = pdf.mean()
    _pdf = pdf.fillna(pdf_mean)
    pdf_std = pdf.std()
    zero_std_labels = _pdf.columns[pdf_std == 0]
    if zero_std_labels.empty:
        return (_pdf - pdf_mean) / pdf_std
    else:
        raise ValueError(f"""Column(s) {zero_std_labels.tolist()} have zero std. dev.!""")


@typechecked
def _prepare_labels_and_warn(label_df: pd.DataFrame, is_binary: bool,
                             label_type: str) -> pd.DataFrame:
    """
    Gets a Pandas DataFrame containing phenotypes (label_df) and whether it is binary or not
    and generates the appropriate warning based on label_type. Returns standardized label DataFrame

    Args:
        label_df: Pandas DataFrame of phenotypes
        is_binary: Whether label_df is binary or not
        label_type: User's choice for label type treatment ('detect', 'binary', or 'quantitative')

    Returns:
        Standardized label Pandas Dataframe
    """
    if label_type == 'detect':
        if is_binary:
            print(
                "The label DataFrame is binary. Reduction/regression for binary phenotypes will be applied.",
            )
        else:
            print(
                "The label DataFrame is quantitative. Reduction/regression for quantitative phenotypes will be applied.",
            )
        return _fill_na_and_standardize(label_df)
    elif label_type == 'quantitative':
        print("Reduction/regression for quantitative phenotypes will be applied.")
        return _fill_na_and_standardize(label_df)
    elif label_type == 'binary':
        if is_binary:
            print("Reduction/regression for binary phenotypes will be applied.")
            return _fill_na_and_standardize(label_df)
        else:
            _check_binary(label_df)
            raise TypeError("Binary label DataFrame expected!", UserWarning)
    else:
        raise ValueError(
            f'label_type should be "quantitative", "binary", or "detect". Found {label_type}.')


@typechecked
def _prepare_covariates(cov_df: pd.DataFrame, label_df: pd.DataFrame,
                        add_intercept: bool) -> pd.DataFrame:
    """
    Standardizes the covariate Pandas DataFrame and adds intercept column to it if indicated

    Args:
        cov_df: Pandas DataFrame of phenotypes
        label_df: Pandas DataFrame of phenotypes
        add_intercept: Whether to add intercept to the covariates Pandas DataFrame

    Returns:
        Standardized covariate Pandas DataFrame with optionally added intercept column
    """
    if cov_df.empty:
        if add_intercept:
            std_cov_df = pd.DataFrame(data=np.ones(label_df.shape[0]),
                                      columns=['intercept'],
                                      index=label_df.index)
        else:
            std_cov_df = cov_df
    else:
        if label_df.shape[0] != cov_df.shape[0]:
            raise ValueError(
                f'cov_df must be either empty of have the same number of rows as label_df ({label_df.shape[0]} != {cov_df.shape[0]}'
            )
        std_cov_df = _fill_na_and_standardize(cov_df)
        if add_intercept:
            std_cov_df.insert(0, 'intercept', 1)

    return std_cov_df


@typechecked()
def _check_model(model_df: DataFrame) -> None:
    """
    Raise an error if model_df does not exist

    Args:
        model_df : Model DataFrame
    """
    if model_df is None:
        raise ValueError(
            'No model DataFrame found! Run fit() or provide a previously made model using set_model_df()'
        )


@typechecked()
def _check_cv(cv_df: DataFrame) -> None:
    """
    Raise an error if cv_df does not exist

    Args:
        cv_df: Cross Validation DataFrame
    """
    if cv_df is None:
        raise ValueError(
            'No cross validation DataFrame found! Run fit() or provide a previously made cv DataFrame using set_cv_df().'
        )


@typechecked
def flatten_prediction_df(blockdf: DataFrame, sample_blocks: Dict[str, Iterable[str]],
                          labeldf: pd.DataFrame) -> pd.DataFrame:
    """
    Transforms a Spark DataFrame containing a block matrix of model predictions to a flat Pandas DataFrame with the same
    shape as labeldf.

    Args:
        blockdf : Spark DataFrame containing a block matrix of predictions for the labels in labeldf
        sample_blocks : Dict containing a mapping of sample_block ID to a list of corresponding sample IDs
        labeldf : Pandas DataFrame containing target labels

    Returns:
        Pandas DataFrame containing prediction values. The shape and order match labeldf such that the
            rows are indexed by sample ID and the columns by label
    """
    sample_block_df = blockdf.sql_ctx \
        .createDataFrame(sample_blocks.items(), ['sample_block', 'sample_ids']) \
        .selectExpr('sample_block', 'posexplode(sample_ids) as (idx, sample_id)')

    flattened_prediction_df = blockdf \
        .selectExpr('sample_block', 'label', 'posexplode(values) as (idx, value)') \
        .join(sample_block_df, ['sample_block', 'idx'], 'inner') \
        .select('sample_id', 'label', 'value')

    pivoted_df = flattened_prediction_df.toPandas() \
        .pivot(index='sample_id', columns='label', values='value') \
        .reindex(index=labeldf.index, columns=labeldf.columns)

    return pivoted_df


@typechecked
def infer_chromosomes(blockdf: DataFrame) -> List[str]:
    """
    Extracts chromosomes from a once- or twice-reduced block DataFrame.

    Args:
        blockdf : Spark DataFrame representing a once- or twice-reduced block matrix.

    Returns:
        List of chromosomes.
    """
    # Regex captures the chromosome name in the header
    # level 1 header: chr_3_block_8_alpha_0_label_sim100
    # level 2 header: chr_3_alpha_0_label_sim100
    chromosomes = [
        r.chromosome for r in blockdf.select(
            f.regexp_extract('header', r"^chr_(.+?)_(alpha|block)", 1).alias(
                'chromosome')).distinct().collect()
    ]
    print(f'Inferred chromosomes: {chromosomes}')
    return chromosomes


def cross_validation(blockdf, modeldf, score_udf, score_key_pattern, alphas, metric):
    """
    Performs cross-validated optimization over the shrinkage hyperparameters alpha, using the models in modeldf and
    the block feature matrix blockdf.

    Args:
        blockdf : Spark DataFrame representing a once- or twice-reduced block matrix.
        modeldf : Spark DataFrame produced by the LogisticRegression or RidgeRegression fit method, representing the
        reducer models fit using the shrinkage parameters alpha.
        score_udf : Pandas UDF used to apply the scoring function to blockdf and modeldf
        score_key_pattern : Pattern of column names used in the groupBy clause before applying score_udf
        alphas : dict of alpha_name : alpha_value
        metric : string specifying either the r2 or the log_loss metric for model scoring

    Returns:
        Spark DataFrame  containing the results of the cross validation procedure.
    """
    alpha_df = blockdf.sql_ctx \
        .createDataFrame([Row(alpha=k, alpha_value=float(v)) for k, v in alphas.items()])
    if metric == 'r2':
        window_spec = Window.partitionBy('label').orderBy(f.desc('r2_mean'), f.desc('alpha_value'))
    elif metric == 'log_loss':
        window_spec = Window.partitionBy('label').orderBy('log_loss_mean', f.desc('alpha_value'))
    else:
        raise ValueError(f'Metric should be either "r2" or "log_loss", found {metric}')

    # Hint a sort-merge join to avoid an automatic broadcast join that may cause an OOM
    return blockdf.drop('header_block', 'sort_key') \
        .join(modeldf.hint('merge'), ['header', 'sample_block'], 'right') \
        .withColumn('label', f.coalesce(f.col('label'), f.col('labels').getItem(0))) \
        .groupBy(score_key_pattern) \
        .apply(score_udf) \
        .join(alpha_df, ['alpha']) \
        .groupBy('label', 'alpha', 'alpha_value').agg(f.mean('score').alias(f'{metric}_mean')) \
        .withColumn('modelRank', f.row_number().over(window_spec)) \
        .filter('modelRank = 1') \
        .drop('modelRank')


def apply_model_df(blockdf, modeldf, cvdf, transform_udf, transform_key_pattern, join_type):
    """
    Applies a model to a reduced block matrix and returns the predictions.

    Args:
        blockdf : Spark DataFrame representing a once- or twice-reduced block matrix.
        modeldf : Spark DataFrame produced by the LogisticRegression or RidgeRegression fit method, representing the
        reducer models fit using the shrinkage parameters alpha.
        cvdf : Spark DataFrame returning the cross validation results
        transform_udf : Pandas UDF used to apply the model
        transform_key_pattern : Pattern of column names used in the groupBy clause
        join_type : Join type for join between blockdf and modeldf. 
    """

    # Hint a sort-merge join to avoid an automatic broadcast join that may cause an OOM
    return blockdf.drop('header_block', 'sort_key') \
        .join(modeldf.drop('header_block').hint('merge'), ['sample_block', 'header'], join_type) \
        .withColumn('label', f.coalesce(f.col('label'), f.col('labels').getItem(0))) \
        .groupBy(transform_key_pattern) \
        .apply(transform_udf) \
        .join(cvdf, ['label', 'alpha'], 'inner')
