import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.decomposition import PCA
from scipy.spatial import distance


def filter_mirnas(cpm_df, threshold=10, min_samples=3):
    """
    Filter miRNAs based on CPM threshold and minimum number of samples.
    cpm_df: DataFrame with miRNAs as rows and samples as columns, containing CPM values.
    threshold: CPM threshold to consider a miRNA as expressed.
    min_samples: Minimum number of samples that must have CPM >= threshold for a miRNA to be retained.
    Returns a filtered DataFrame with only the miRNAs that meet the criteria.
    """
    # Create a boolean mask where True if cpm >= threshold
    mask = cpm_df >= threshold
    # Sum across rows to count how many samples have cpm >= threshold for each miRNA
    count_above_threshold = mask.sum(axis=1)
    # Filter miRNAs that have at least min_samples with cpm >= threshold
    filtered_cpm_df = cpm_df[count_above_threshold >= min_samples]
    return filtered_cpm_df


def run_pca_and_plot_samples(raw_counts_df, samples_df, group='IP-Nat'):
    """"
    This function performs PCA on the log2-transformed raw counts and creates a scatter plot of the first two principal components.
    Each point represents a sample, colored by the specified group (e.g., 'IP-Nat'). 
    Sample names are added as labels to help identify any potential outliers or clustering patterns.

    """
    # 1. Prepare the data for PCA
    df_pca_input = np.log2(raw_counts_df + 1).T 

    # 2. Run PCA
    pca = PCA(n_components=2)
    components = pca.fit_transform(df_pca_input)
    exp_var = pca.explained_variance_ratio_ * 100

    # 3. Create a DataFrame for plotting
    pca_df = pd.DataFrame(
        data = components, 
        columns = ['PC1', 'PC2'],
        index = raw_counts_df.columns
    )

    # Add the group information for coloring
    pca_df['Condition'] = samples_df.set_index('Sample').loc[pca_df.index][group]

    # 4. Plot
    plt.figure(figsize=(10, 7))
    sns.scatterplot(
        x='PC1', y='PC2', 
        hue='Condition', 
        data=pca_df, 
        s=100, 
        style='Condition'
    )

    # Add labels with sample names to identify outliers
    for i in range(pca_df.shape[0]):
        plt.text(
            pca_df.PC1[i]+0.2, 
            pca_df.PC2[i]+0.2, 
            pca_df.index[i], 
            fontsize=9, 
            alpha=0.7
        )

    plt.title('PCA Plot - Checking Group Separation', fontsize=15)
    plt.xlabel(f'PC1 ({exp_var[0]:.1f}% variance)')
    plt.ylabel(f'PC2 ({exp_var[1]:.1f}% variance)')
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.show()


def get_correlation_matrix(raw_counts_df, save_name=None):
    """
    This function calculates the correlation of the samples, which gives a measure of how similar the samples are to each other.
    A correlation of 1 means the samples are identical, while a correlation less than 0.8 may indicate 
    a potential outlier or a sample that is not behaving like the others.

    The heatmap visualizes these correlations, with lighter colors indicating higher correlation, and darker colors indicating lower correlation.
    Colors in vibrant red across all the samples would indicate potential outliers or samples that are not behaving like the others.

    params:
    :raw_counts_df: DataFrame of raw counts with miRNAs as rows and samples as columns
    :save_name: Optional string for saving the heatmap figure (e.g., 'correlation_matrix.svg')
    """
    # Calculate the correlation matrix using log2-transformed counts
    corr_matrix = np.log2(raw_counts_df + 1).corr()

    # print heatmap of correlation matrix

    plt.figure(figsize=(10, 8))
    sns.heatmap(corr_matrix, annot=False, cmap='Reds_r', vmin=0.8, vmax=1.0)
    plt.title('Correlation Matrix of Samples')
    if save_name:
        plt.savefig(save_name)
    plt.show()
    return corr_matrix


def run_outlier_detection(cpm_df, sample_list, threshold=1.5):
    """
    Detects outliers by measuring Euclidean distance from the group centroid. This
    function will take the log2-transformed CPM values, calculate the centroid (average profile), 
    and then compute the distance of each sample from this centroid. 
    Samples that are more than threshold standard deviations away from the mean distance will be flagged as outliers.
    """
    data = np.log2(cpm_df[sample_list] + 1).T
    centroid = data.mean(axis=0)
    
    # Distance of each sample from the 'average' sample
    distances = data.apply(lambda row: distance.euclidean(row, centroid), axis=1)
    z_scores = (distances - distances.mean()) / distances.std()
    
    outlier_report = pd.DataFrame({
        'Euclidean_Distance': distances,
        'Z_Score': z_scores,
        'Is_Outlier': z_scores > threshold # Flag if > threshold standard deviations away
    })
    return outlier_report
