"""
MDTerp.init_analysis.py – Initial MDTerp round for discarding irrelevant features from analysis for the forward feature selection in MDTerp.final_analysis.py.
"""
import numpy as np
import os
import sklearn.metrics as met
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as lda
from sklearn.linear_model import Ridge
from typing import Tuple

def similarity_kernel(data: np.ndarray, kernel_width: float = 1.0) -> np.ndarray:
    """
    Function for computing similarity∈[0,1] of a perturbed sample with respect to the original sample using LDA transformed distance.

    Args:
        data (np.ndarray): LDA transformed data.
        kernel_width (float): Width of the similarity kernel (Default: 1.0).

    Returns:
        np.ndarray: Similarity∈[0,1] of neighborhood.
    """
    distances = met.pairwise_distances(data,data[0].reshape(1, -1),metric='euclidean').ravel()
    return np.sqrt(np.exp(-(distances ** 2) / kernel_width ** 2))

def SGDreg(data: np.ndarray, labels: np.ndarray, seed: int, alpha: float = 1.0) -> Tuple[np.ndarray, float]:
    """
    Function for implementing linear regression using stochastic gradient descent.

    Args:
        data (np.ndarray): Numpy 2D array containing similarity weighted training data for the black-box model. Samples along rows and features along columns.
        labels (np.ndarray): Numpy array containing metastable state prediction probabilities for a perturbed neighborhood corresponding to a specific sample. Includes the state for which the original sample has the highest probability.
        seed (int): Random seed.
        alpha (float): L2 norm of Ridge regression (Default: 1.0).
        
    Returns:
        np.ndarray: Numpy array with coefficients of all the features of the fitted linear model.
        float: Intercept of the fitted linear model.
    """
    clf = Ridge(alpha, random_state = seed, solver = 'saga')
    clf.fit(data,labels.ravel())
    coefficients = clf.coef_
    intercept = clf.intercept_
    return coefficients, intercept

def selection(coefficients: np.ndarray, threshold: float) -> np.ndarray:
    """
    Function for discarding irrelevant features with small absolute values of the coefficients from the fitted linear model.

    Args:
        coefficients (np.ndarray): Numpy array with coefficients of all the features of the fitted linear model.
        threshold (float): Hyperparameter controlling how many features are discarded.
        
    Returns:
        np.ndarray: Numpy array containing indices of the selected features for further analysis/forward feature selection.
    """
    coefficients_abs = np.absolute(coefficients)
    selected_features = []
    coverage = 0
    for i in range(coefficients_abs.shape[0]):
        if i==threshold:
            break
        coverage = coverage+np.sort(coefficients_abs)[::-1][i]/np.sum(coefficients_abs)
        selected_features.append(np.argsort(coefficients_abs)[::-1][i])
    return selected_features    
    
def init_model(neighborhood_data: np.ndarray, pred_proba: np.ndarray, cutoff: float, given_indices: np.ndarray, seed: int) -> list:
    """
    Function fitting initial linear model for discarding irrelevant features and choosing promising features for detailed analysis.

    Args:
        neighborhood_data (np.ndarray): Perturbed data generated by MDTerp.neighborhood.py.
        pred_proba (np.ndarray): Metastable state probabilities obtained from the black-box.
        cutoff (float): Hyperparameter controlling how many features are discarded.
        given_indices (np.ndarray): Indices of the features to perform first round of MDTerp on.
        seed (int): Random seed.
        
    Returns:
        np.ndarray: List of three lists containing indices of the selected numeric, angular, sine/cosine features respectively for further analysis/forward feature selection.
    """
    explain_class = np.argmax(pred_proba[0,:])

    target = pred_proba[:,explain_class]

    threshold, upper, lower = 0.5, 1, 0
    target_binarized = np.where(target>threshold, upper, lower)

    clf = lda()
    clf.fit(neighborhood_data,target_binarized)
    projected_data = clf.transform(neighborhood_data)
    weights = similarity_kernel(projected_data.reshape(-1,1), 1)

    
    predict_proba = pred_proba[:,explain_class]
    data = neighborhood_data*(weights**0.5).reshape(-1,1)
    labels = target.reshape(-1,1)*(weights.reshape(-1,1)**0.5)
    
    coefficients_selection, intercept_selection = SGDreg(data, labels, seed)
    coefficients_selection = coefficients_selection/np.sum(np.absolute(coefficients_selection))
    selected_features = selection(coefficients_selection, cutoff)

    selected_features_lst = [[], [], []]
    
    for i in range(len(selected_features)):
        if selected_features[i]<given_indices[0].shape[0]:
            selected_features_lst[0].append(given_indices[0][selected_features[i]])
        else:
            if selected_features[i]<given_indices[0].shape[0] + given_indices[1].shape[0]:
                selected_features_lst[1].append(given_indices[1][selected_features[i] - given_indices[0].shape[0]])
            else:
                selected_features_lst[2].append(given_indices[2][selected_features[i] - given_indices[0].shape[0] - given_indices[1].shape[0], :])

    return selected_features_lst
