import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


def prepare_core_data(cpm_path, counts_path, metadata_path):
    """
    Standardizes indices and column names for all core dataframes.
    """
    # Load
    cpm_df = pd.read_csv(cpm_path)
    raw_df = pd.read_csv(counts_path)
    meta_df = pd.read_csv(metadata_path)
    
    # Standardize CPM
    cpm_df.rename(columns={'Name': 'miRNA'}, inplace=True)
    cpm_df.set_index('miRNA', inplace=True)
    cpm_df.columns = cpm_df.columns.str.replace(" - CPM", "")
    
    # Standardize Raw Counts
    raw_df.rename(columns={'Name': 'miRNA'}, inplace=True)
    raw_df.set_index('miRNA', inplace=True)
    raw_df.columns = raw_df.columns.str.replace(" - Total counts", "")
    raw_df = raw_df.select_dtypes(include=[np.number]) # Keep only numeric
    
    return cpm_df, raw_df, meta_df


from kneed import KneeLocator
from scipy.spatial import distance
import scipy.stats as stats

def get_elbow_threshold(cpm_df, plot=False):
    """
    Finds the elbow point to distinguish signal from noise.
    """
    avg_cpm = cpm_df.select_dtypes(include=[np.number]).mean(axis=1)
    avg_cpm_sorted = avg_cpm.sort_values(ascending=False).reset_index(drop=True)
    
    # KneeLocator finds the 'curve' point
    knee = KneeLocator(range(len(avg_cpm_sorted)), avg_cpm_sorted, 
                       curve='convex', direction='decreasing')
    
    if plot:
        plt.plot(avg_cpm_sorted.values)
        plt.axhline(y=knee.knee_y, color='r', linestyle='--')
        plt.title(f"Elbow Threshold Identified: {knee.knee_y:.2f} CPM")
        plt.show()
        
    return knee.knee_y

def calculate_group_consistency(cpm_df, sample_list, group_name="Group"):
    """
    Calculates how well replicates correlate with each other.
    """
    # Use log2 to prevent high-expressors from biasing correlation
    corr_matrix = np.log2(cpm_df[sample_list] + 1).corr()
    avg_corr = corr_matrix.values.mean()
    
    print(f"Average internal correlation for {group_name}: {avg_corr:.4f}")
    return avg_corr


def run_deseq2_pipeline(counts_df, group_labels, group_1='group_1', group_2='group_2',
                        min_counts=10, min_samples=3):
    """
    counts_df: Must be RAW INTEGERS. Rows=miRNA, Cols=Samples.
    group_labels: Must match the length of counts_df.columns.
    group_1: Label for the numerator group in the contrast (default: 'group_1').
    group_2: Label for the denominator group in the contrast (default: 'group_2').
    """
    from anndata import AnnData
    from pydeseq2.dds import DeseqDataSet, DefaultInference
    from pydeseq2.ds import DeseqStats

    # --- SANITY CHECKS ---
    # 1. Ensure labels match columns
    if len(group_labels) != len(counts_df.columns):
        raise ValueError(f"Label length ({len(group_labels)}) doesn't match column count ({len(counts_df.columns)})")
    
    # 2. Force Integers (DESeq2 will crash on floats/CPM)
    # This rounds any normalization artifacts and casts to int
    counts_df = counts_df.round().astype(int)

    # --- FILTERING ---
    # Keep miRNA if it has >= 10 reads in at least 3 mice
    keep = (counts_df >= min_counts).sum(axis=1) >= min_samples
    filtered_counts = counts_df[keep].T # Transpose for AnnData (Samples are now rows)
    
    print(f"DESeq2 Input: {filtered_counts.shape[1]} miRNAs pass the {min_counts} count filter.")

    # --- EXECUTION ---
    adata = AnnData(X=filtered_counts, 
                    obs=pd.DataFrame({'condition': group_labels}, 
                    index=filtered_counts.index))
    
    dds = DeseqDataSet(adata=adata, design_factors='condition', refit_cooks=True, inference=DefaultInference(n_cpus=8))
    dds.deseq2()
    
    # Contrast: [Column_Name, Numerator, Denominator]
    stat_res = DeseqStats(dds, contrast=["condition", group_1, group_2], inference=DefaultInference(n_cpus=8))
    stat_res.summary()
    
    return stat_res.results_df


def get_dds(andata):
        """
        This function will return the DESeq2 object dds needed to perfome the DEA
        """
        from pydeseq2.dds import DeseqDataSet, DefaultInference

        # Build DESeq2 object
        
        dds = DeseqDataSet(
            adata=andata,
            design_factors='condition',
            refit_cooks=True,
            inference=DefaultInference(n_cpus=8),
        )
        dds.deseq2()
        return dds


def get_DEA(dds, group_1='young', group_2='old', save=True):
        from pydeseq2.dds import DefaultInference
        from pydeseq2.ds import DeseqStats

        comparison = f'{group_1}.vs.{group_2}'
        stat_res = DeseqStats(
            dds,
            contrast=["condition", group_1, group_2],
            inference=DefaultInference(n_cpus=8)
        )
        stat_res.summary()
        results_df = stat_res.results_df
        return results_df


def get_exclusion_lists(cpm_df, group_1_samples, group_2_samples, threshold=0):
    """
    Finds miRNAs exclusively expressed in one group or the other.
    """
    # Expressed in at least one sample of the group
    group_1_active = cpm_df.index[(cpm_df[group_1_samples] > threshold).any(axis=1)]
    group_2_active = cpm_df.index[(cpm_df[group_2_samples] > threshold).any(axis=1)]
    
    unique_to_group_1 = set(group_1_active) - set(group_2_active)
    unique_to_group_2 = set(group_2_active) - set(group_1_active)
    
    return list(unique_to_group_1), list(unique_to_group_2)


def plot_personalize_bars(dea, title="", thr_significance=0.05, thr_stat=0.5, top=20, max_stat=None,
                          title_fontsize=14, figsize=(5, 5), dpi=700, filename=None):
    import decoupler as dc
    # Generate the volcano plot
    fig = dc.pl.volcano(
        data = dea,
        x='stat',
        y='padj',
        top=top,
        thr_sign=thr_significance,
        thr_stat=thr_stat,
        max_stat=max_stat,
        return_fig=True,
        figsize=figsize,
    )
    
    # Access the Axes object from the figure
    ax = fig.gca()
    
    # Add a title and remove the outer black line
    ax.set_title(title, fontsize=title_fontsize)
    for spine in ax.spines.values():
        spine.set_visible(False)  # Hide all outer spines
    
    return fig, ax
