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


def plot_barplot_difference_mirs(
    cpm_df,
    samples_df,
    group="IP-Nat",
    threshold=10,
    save_name=None,
):
    """
    Calculate average number of expressed miRNAs per group and plot a barplot.

    cpm_df: miRNAs x samples (columns are sample IDs)
    samples_df: must contain "Sample" and the grouping column
    """
    if group not in samples_df.columns:
        raise ValueError(f'Group column "{group}" not found in samples_df')

    common_samples = cpm_df.columns.intersection(samples_df["Sample"])
    if len(common_samples) == 0:
        raise ValueError("No common samples found between cpm_df and samples_df")

    cpm = cpm_df.loc[:, common_samples]
    meta = samples_df.loc[samples_df["Sample"].isin(common_samples), ["Sample", group]]

    groups = meta[group].dropna().unique()
    group_counts_d = {}
    expressed_means = []

    for g in groups:
        s = meta.loc[meta[group] == g, "Sample"]
        group_expressed = (cpm.loc[:, s] > threshold).sum(axis=0)
        group_counts_d[g] = group_expressed
        expressed_means.append(group_expressed.mean())

    plt.figure(figsize=(7, 5))
    plt.bar(groups, expressed_means, color=["blue", "orange"][: len(groups)])
    plt.ylabel(f"Average # expressed miRNAs (>{threshold} CPM)")
    plt.title(f"Average # expressed miRNAs by {group}")
    if save_name:
        plt.savefig(save_name, bbox_inches="tight", dpi=300)
    plt.show()

    return group_counts_d


def plot_boxplot_difference_mirs(
    cpm_df,
    samples_df,
    group="IP-Nat",
    threshold=10,
    pvalue=0.05,
    save_name=None,
):
    """
    Calculate number of expressed miRNAs per sample within each group and plot a boxplot.
    Also prints t-test between the first two groups found.
    """
    if group not in samples_df.columns:
        raise ValueError(f'Group column "{group}" not found in samples_df')

    common_samples = cpm_df.columns.intersection(samples_df["Sample"])
    if len(common_samples) == 0:
        raise ValueError("No common samples found between cpm_df and samples_df")

    cpm = cpm_df.loc[:, common_samples]
    meta = samples_df.loc[samples_df["Sample"].isin(common_samples), ["Sample", group]]

    groups = meta[group].dropna().unique()
    group_counts_d = {}

    for g in groups:
        s = meta.loc[meta[group] == g, "Sample"]
        group_counts_d[g] = (cpm.loc[:, s] > threshold).sum(axis=0)

    # t-test (only meaningful for 2 groups; kept consistent with your original behavior)
    t_stat, p_value = ttest_ind(*list(group_counts_d.values())[:2])
    print(f"T-statistic: {t_stat}, P-value: {p_value}")

    plt.figure(figsize=(8, 6))
    sns.boxplot(data=pd.DataFrame(group_counts_d), palette="Set2")
    plt.title(f"Distribution of expressed miRNAs in each {group} group")
    plt.ylabel("Number of expressed miRNAs")
    plt.xlabel(group)
    # add the bar with the * if significant
    if p_value < pvalue:
        ax = plt.gca()

        # positions of the two boxes (works for 2 groups)
        x1, x2 = 0, 1

        # height for the bracket
        data_max = max(pd.concat(group_counts_d.values()))
        y = data_max * 1.05
        h = data_max * 0.05  # bracket height

        # draw bracket
        ax.plot([x1, x1, x2, x2], [y, y+h, y+h, y], lw=1.5, c="black")

        # add asterisk
        ax.text((x1 + x2) / 2, y + h, f"*{f'p = {p_value:.3e}'}", 
                ha="center", va="bottom", fontsize=16)

    if save_name:
        plt.savefig(save_name, bbox_inches="tight", dpi=300)
    plt.show()

    return group_counts_d


def plot_burden_and_richness(analysis_df, title="Global miRNA Profile", condition='Condition',
                             burden='Total_burden', miR_count='miR_count', threshold=0.05):
    """
    Creates a two-panel figure showing Total Signal (Burden) and Unique miR counts (Richness).

    Parameters:
    -----------
    analysis_df : pd.DataFrame
        DataFrame returned by calculate_burden_richness.
        Columns expected: ['Total_Burden', 'MiR_Count', 'Condition']
        Index: Sample IDs.
    title : str
        Main title for the figure.
    """
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    groups = analysis_df[condition].unique()
    gs_burden = []
    gs_richness = []
    for group in groups:
        if group not in analysis_df[condition].values:
            raise ValueError(f"Group '{group}' not found in the '{condition}' column of the DataFrame.")
        gs_burden.append(analysis_df[analysis_df[condition] == group][burden])
        gs_richness.append(analysis_df[analysis_df[condition] == group][miR_count])

    t_stat, p_value = ttest_ind(*gs_burden)
    if p_value < 0.05:
        y = analysis_df[burden].max() * 1.05
        h = analysis_df[burden].max() * 0.05
        axes[0].plot([0, 0, 1, 1], [y, y+h, y+h, y], lw=1.5, c='black')
        axes[0].text(0.5, y+h, '*', ha='center', va='bottom', fontsize=16)
    title_B = "Total Burden (Sum CPM)"
    if p_value < threshold:
        title_B += f" (p={p_value:.4f})"
    sns.boxplot(x=condition, y=burden, data=analysis_df, palette='Set2', ax=axes[0])
    sns.stripplot(x=condition, y=burden, data=analysis_df, color='black', alpha=0.5, ax=axes[0])
    axes[0].set_title(title_B)

    t_stat, p_value = ttest_ind(*gs_richness)
    if p_value < 0.05:
        y = analysis_df[miR_count].max() * 1.05
        h = analysis_df[miR_count].max() * 0.05
        axes[1].plot([0, 0, 1, 1], [y, y+h, y+h, y], lw=1.5, c='black')
        axes[1].text(0.5, y+h, '*', ha='center', va='bottom', fontsize=16)
    title_R = "Richness (Unique miRNAs)"
    if p_value < threshold:
        title_R += f" (p={p_value:.4f})"
    sns.boxplot(x=condition, y=miR_count, data=analysis_df, palette='Pastel1', ax=axes[1])
    sns.stripplot(x=condition, y=miR_count, data=analysis_df, color='black', alpha=0.5, ax=axes[1])
    axes[1].set_title(title_R)

    plt.suptitle(title, fontsize=16)
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    plt.show()


def plot_global_enrichment(group_1_enrichment, group_2_enrichment, group="Global", threshold=0.05,
                           group_1_label="Group1", group_2_label="Group2"):
    # 1. Calculate the Global Average Enrichment per Mouse
    # We take the mean across all miRNAs for each mouse column
    group_1_mean_enrichment = group_1_enrichment.mean()
    group_2_mean_enrichment = group_2_enrichment.mean()

    # 2. Prepare data for plotting
    plot_data = pd.DataFrame({
        'Average_Enrichment': pd.concat([group_1_mean_enrichment, group_2_mean_enrichment]),
        'Group': [group_1_label] * len(group_1_mean_enrichment) + [group_2_label] * len(group_2_mean_enrichment)
    })

    # 3. Perform T-test
    t_stat, p_val = stats.ttest_ind(group_1_mean_enrichment, group_2_mean_enrichment)

    print(f"--- {group} Enrichment Analysis ---")
    print(plot_data.groupby('Group')['Average_Enrichment'].agg(['mean', 'std', 'count']))
    print(f"\nT-test Result: t = {t_stat:.4f}, p-value = {p_val:.4f}")

    # 4. Create the Boxplot
    plt.figure(figsize=(6, 7))
    sns.set_style("whitegrid")

    # Boxplot for the distribution
    sns.boxplot(x='Group', y='Average_Enrichment', data=plot_data, palette='Set1', width=0.5)

    # Stripplot to show individual mice
    sns.stripplot(x='Group', y='Average_Enrichment', data=plot_data, color='black', alpha=0.6, size=8)

    plt.title(f'{group} $o^8$ Enrichment Burden\n(Paired $\\log_2(IP/NAT)$ per mouse)', fontsize=14)
    plt.ylabel('Mean Enrichment Score', fontsize=12)
    plt.xlabel('Group', fontsize=12)
    # add label with the p-value
    plt.text(0.5, plot_data['Average_Enrichment'].max() + 0.1, f'p = {p_val:.4f}', ha='center', va='bottom', color='black', fontsize=12)
    # Optional: Add significance star if p < threshold
    if p_val < threshold:
        x1, x2 = 0, 1
        y, h = plot_data['Average_Enrichment'].max() + 0.1, 0.1
        plt.plot([x1, x1, x2, x2], [y, y+h, y+h, y], lw=1.5, c='black')
        plt.text((x1+x2)*.5, y+h, "*", ha='center', va='bottom', color='black', fontsize=16)
    
    plt.tight_layout()
    plt.show()

    return plot_data, t_stat, p_val


def plot_kde(group_1_enrichment, group_2_enrichment, group='Global',
             group_1_label="Group1", group_2_label="Group2"):
    #1. Reshape the data to "Long Form"
    # We melt the data so every row is one miRNA in one mouse
    group_1_melted = group_1_enrichment.melt(var_name='Mouse', value_name='Enrichment')
    group_1_melted['Group'] = group_1_label

    group_2_melted = group_2_enrichment.melt(var_name='Mouse', value_name='Enrichment_G2')
    group_2_melted['Group'] = group_2_label

    # Combine them
    all_data_points = pd.concat([
        group_1_melted[['Enrichment', 'Group']], 
        group_2_melted.rename(columns={'Enrichment_G2': 'Enrichment'})[['Enrichment', 'Group']]
    ])

    # 2. Statistics: Kolmogorov-Smirnov Test
    # This tests if the two curves (distributions) are significantly different
    ks_stat, ks_p_val = stats.ks_2samp(
        group_1_melted['Enrichment'].dropna(), 
        group_2_melted['Enrichment_G2'].dropna()
    )

    print(f"--- {group} Distribution Analysis ---")
    print(f"KS Statistic: {ks_stat:.4f}")
    print(f"KS P-value: {ks_p_val:e}") # Often very small

    # 3. Visualization: Ridge Plot / Density Plot
    plt.figure(figsize=(10, 6))

    sns.kdeplot(data=all_data_points, x='Enrichment', hue='Group', fill=True, common_norm=False, palette='Set1', alpha=0.5)

    plt.axvline(group_1_melted['Enrichment'].median(), color='red', linestyle='--', label=f'{group_1_label} Median')
    plt.axvline(group_2_melted['Enrichment_G2'].median(), color='blue', linestyle='--', label=f'{group_2_label} Median')

    plt.title(f'{group} Shift in miRNA Oxidation Propensity\n(Distribution of all miRNA enrichment scores)', fontsize=14)
    plt.xlabel('Enrichment Score: $Log_2(IP/NAT)$')
    plt.ylabel('Density')
    plt.legend()
    plt.show()

    print(f"{group_1_label} Median Enrichment: {group_1_melted['Enrichment'].median():.4f}")
    print(f"{group_2_label} Median Enrichment: {group_2_melted['Enrichment_G2'].median():.4f}")

    shift = group_2_melted['Enrichment_G2'].median() - group_1_melted['Enrichment'].median()
    percentage_increase = (2**shift - 1) * 100
    print(f"The {group_2_label} group increases the global proportion of oxidation by {percentage_increase:.2f}%")

    return all_data_points, {"KS-statistic": ks_stat,
                                "KS-p-value": ks_p_val,
                                "Shift in median enrichment (log2 scale)": shift,
                                "Group 1 median": group_1_melted['Enrichment'].median(),
                                "Group 2 median": group_2_melted['Enrichment_G2'].median(),
                                "Group 2 increase global proportion": percentage_increase*100}


def ECDF_plot(clean_enrichment, group_1_individuals, group_2_individuals, group_1_name, 
              group_2_name, title="Cumulative Distribution of Oxidation Enrichment", ignore_missing=False):
    """
        This function plots the ECDF (Empirical Cumulative Distribution Function) 
        for two groups of individuals O8G enrichment. Calculated as the log2 fold change
        between the IP and NAT conditions.

        parameters:
            clean_enrichment: pd.DataFrame
                A DataFrame containing the enrichment scores (log2 IP/NAT) for all miRNAs and all individuals. Rows = miRNAs, Columns = Individual Mouse IDs.
            group_1_individuals: list
                A list of strings containing the Sample IDs for the first group (e.g., Vehicle mice).
            group_2_individuals: list
                A list of strings containing the Sample IDs for the second group (e.g., KO mice).
            group_1_name: str
                A string representing the name of the first group (e.g., "Vehicle").
            group_2_name: str
                A string representing the name of the second group (e.g., "KO").
            title: str
                The title for the plot.
            ignore_missing: bool
                If True, the function will ignore any individuals in group_1_individuals or group_2_individuals
                that are not in clean_enrichment.columns.

    """
    # check that all individuals are in clean_enrichment
    all_individuals = set(group_1_individuals + group_2_individuals)
    missing_individuals = all_individuals - set(clean_enrichment.columns)
    if missing_individuals:
        if ignore_missing:
            print(f"Warning: The following individuals are missing from the enrichment DataFrame and will be ignored: {missing_individuals}")
            group_1_individuals = [ind for ind in group_1_individuals if ind in clean_enrichment.columns]
            group_2_individuals = [ind for ind in group_2_individuals if ind in clean_enrichment.columns]
        else:
            raise ValueError(f"The following individuals are missing from the enrichment DataFrame: {missing_individuals}")
    # ECDF plot - This is often better for showing the 3% shift mentioned in the KS-test
    plt.figure(figsize=(8, 6))
    sns.ecdfplot(data=clean_enrichment[group_1_individuals].melt(), x='value', label=group_1_name, color='red')
    sns.ecdfplot(data=clean_enrichment[group_2_individuals].melt(), x='value', label=group_2_name, color='blue')
    plt.title(title)
    plt.legend()
    plt.show()


def plot_positional_vulnerability(mismatch_df, group_1_samples, group_2_samples,
                                  condition="Exercise", group_1_label="Group1", group_2_label="Group2"):
    """
    Calculates the G>T mutation rate for each specific sequence position and plots the distribution.
    
    Parameters:
    -----------
    mismatch_df : pd.DataFrame
        The raw mismatch data containing 'miRNA name' and 'pos:mut' columns.
    group_1_samples, group_2_samples : lists
        Lists of Sample IDs.
    """
    # Ensure column name is consistent
    if mismatch_df.columns[0] != 'miRNA name':
        # check if index is miRNA name and reset if needed
        if mismatch_df.index.name == 'miRNA name':
            mismatch_df = mismatch_df.reset_index()
        else:            
            raise ValueError("Expected 'miRNA name' as the first column or index. Please check your DataFrame.")
    df = mismatch_df.rename(columns={'miRNA name': 'miRNA'})
    # if type of group_1_samples and group_2_samples are lists, sum
    if isinstance(group_1_samples, list) and isinstance(group_2_samples, list):
        all_samples = group_1_samples + group_2_samples
    else:
        raise ValueError("group_1_samples and group_2_samples should be lists of sample IDs.")

    # 1. Calculate Denominator (Total Reads per miRNA)
    total_reads = df.groupby('miRNA')[all_samples].sum() + 1 # +1 to avoid div zero
    
    # 2. Isolate only the rows containing G>T mutations
    gt_df = df[df['pos:mut'].str.contains('G>T', na=False)].copy()
    
    # 3. Extract the numeric position specifically for these rows
    def get_exact_pos(val):
        parts = val.split(',')
        for p in parts:
            if 'G>T' in p:
                num = ''.join(filter(str.isdigit, p.split(':')[0]))
                return int(num) if num else np.nan
        return np.nan
        
    gt_df['position'] = gt_df['pos:mut'].apply(get_exact_pos)
    gt_df = gt_df.dropna(subset=['position'])

    # 4. Merge total reads into the mutation dataframe
    # This attaches the 'Total' to every instance of that miRNA
    merged_df = gt_df.merge(total_reads, on='miRNA', suffixes=('', '_total'))

    # 5. Calculate the rate: Specific Row Counts / Total miRNA Counts
    rate_results = []
    for col in all_samples:
        # Rate = Count of specific G>T / (Total reads for that miR + 1)
        merged_df[f'{col}_rate'] = merged_df[col] / (merged_df[f'{col}_total'] + 1)

    # 6. Group by Position to get the mean rate across all miRNAs
    rate_cols = [f'{col}_rate' for col in all_samples]
    
    # 7. Group by Position to get the average rate across all miRNAs at position X
    pos_dist = merged_df.groupby('position')[rate_cols].mean()
    # Rename columns back to original sample IDs for easier plotting
    pos_dist.columns = all_samples
    
    # 8. Plotting
    plt.figure(figsize=(12, 5))
    sns.lineplot(data=pos_dist[group_1_samples].mean(axis=1), label=group_1_label, color='green', marker='o')
    sns.lineplot(data=pos_dist[group_2_samples].mean(axis=1), label=group_2_label, color='purple', marker='o')
    
    plt.title(f'Average G>T Mutation Rate by Base Position ({condition})', fontsize=14)
    plt.xlabel('Base Position (5\' -> 3\')', fontsize=12)
    plt.ylabel('Mutation Rate (G>T / Total Reads)', fontsize=12)
    plt.xticks(np.arange(1, 26, 1)) # Force integer ticks for positions 1-25
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()
    
    return pos_dist # this is the dataframe with average G>T mutation rates by position for all samples.


def plot_differential_heatmap(enrichment_df, target_mirs, group_1_individuals, group_2_individuals,
                              title="Selective miRNA Oxidation Targets"):




    """
    Generates a heatmap showing the enrichment scores of specific miRNAs across all samples,
    grouped by genotype.

    Parameters:
    -----------
    enrichment_df : pd.DataFrame
        The matrix of paired enrichment scores (log2 IP/NAT). 
        Rows = miRNA names, Columns = Mouse IDs.
        Expected Data: Numeric floats, typically ranging from -5 to 5.
    target_mirs : list
        A list of strings containing the names of miRNAs that passed the p-value threshold (e.g., p < 0.05).
        Example: ['mmu-miR-1a-3p', 'mmu-miR-133a-3p']
    group_1_individuals : list
        List of strings for group 1 Mouse IDs to be plotted first.
    group_2_individuals : list
        List of strings for group 2 Mouse IDs to be plotted second.
    title : str
        The title for the heatmap.

    Example of Data Expected:
    -------------------------
    enrichment_df:
                    Mouse_1  Mouse_2  Mouse_3 (KO)  Mouse_4 (KO)
    mmu-miR-1       0.12     -0.05    1.45          2.10
    mmu-miR-133     -0.10    0.02     0.88          1.25
    """
    
    # 1. Filter the matrix for only the targets of interest
    plot_matrix = enrichment_df.loc[target_mirs]
    
    # 2. Reorder columns so group 1 mice are on the left and group 2 are on the right
    plot_matrix = plot_matrix[list(group_1_individuals) + list(group_2_individuals)]
    
    # 3. Plotting
    plt.figure(figsize=(10, len(target_mirs) * 0.4 + 2)) # Dynamic height based on number of miRs
    
    # Use 'coolwarm' map: Blue = Unoxidized/Native, Red = Highly Oxidized/Enriched
    # Center the colormap at 0
    sns.heatmap(
        plot_matrix, 
        cmap='coolwarm', 
        center=0, 
        annot=True, 
        fmt=".2f",
        cbar_kws={'label': 'Enrichment Score $Log_2(IP/NAT)$'}
    )
    
    # Add a vertical line to visually separate genotypes
    plt.axvline(len(group_1_individuals), color='black', lw=3)
    
    plt.title(title, fontsize=16)
    plt.xlabel('Individual Mice (Paired Fractions)', fontsize=12)
    plt.ylabel('miRNA Species', fontsize=12)
    
    plt.tight_layout()
    plt.show()


def plot_pca_separation(counts_df, group_1_samples, title='PCA - Group Separation',
                        group_1_label="Group1", group_2_label="Group2"):
    """
    Generates a PCA plot to visualize if samples cluster by Genotype.

    Parameters:
    -----------
    counts_df : pd.DataFrame
        DataFrame of log2-transformed or raw counts. Rows=miRNAs, Cols=Samples.
        Example: 2000 rows x 10 columns.
    group_1_samples : list
        List of strings containing Sample IDs belonging to the first group.
        Example: ['31136-001', '31136-002']
    title : str
        The title for the plot.
    """
    # Logic from run_pca_and_plot_ve_ko in functions_oxy.py
    df_pca_input = np.log2(counts_df + 1).T 
    pca = PCA(n_components=2)
    components = pca.fit_transform(df_pca_input)
    exp_var = pca.explained_variance_ratio_ * 100

    pca_df = pd.DataFrame(data=components, columns=['PC1', 'PC2'], index=counts_df.columns)
    pca_df['Condition'] = [group_1_label if col in group_1_samples else group_2_label for col in pca_df.index]

    plt.figure(figsize=(8, 6))
    sns.scatterplot(x='PC1', y='PC2', hue='Condition', data=pca_df, s=100)
    
    for i in range(pca_df.shape[0]):
        plt.text(pca_df.PC1[i]+0.1, pca_df.PC2[i]+0.1, pca_df.index[i], fontsize=8)

    plt.title(title)
    plt.xlabel(f'PC1 ({exp_var[0]:.1f}%)')
    plt.ylabel(f'PC2 ({exp_var[1]:.1f}%)')
    plt.show()


def plot_enrichment_distribution(group_1_enrichment, group_2_enrichment, condition='Exercise',
                                 group_1_label="Group1", group_2_label="Group2"):
    """
    Generates a KDE Density plot and ECDF to visualize the global shift in oxidation propensity.

    Parameters:
    -----------
    group_1_enrichment : pd.DataFrame
        Enrichment scores (log2 IP/NAT) for group 1 mice. Rows=miRNAs, Cols=Mice.
        Example: 500 rows x 5 columns.
    group_2_enrichment : pd.DataFrame
        Enrichment scores (log2 IP/NAT) for group 2 mice. Rows=miRNAs, Cols=Mice.
    condition : str
        The experimental context (e.g., 'Exercise').
    """
    # Logic from plot_kde in functions_oxy.py
    g1_points = group_1_enrichment.values.flatten()
    g2_points = group_2_enrichment.values.flatten()
    
    plt.figure(figsize=(10, 6))
    sns.kdeplot(g1_points, label=group_1_label, fill=True, color='red', alpha=0.3)
    sns.kdeplot(g2_points, label=group_2_label, fill=True, color='blue', alpha=0.3)
    
    plt.axvline(np.nanmedian(g1_points), color='red', linestyle='--')
    plt.axvline(np.nanmedian(g2_points), color='blue', linestyle='--')
    
    plt.title(f'Oxidation Propensity Shift: {condition}')
    plt.xlabel('Enrichment Score: $Log_2(IP/NAT)$')
    plt.ylabel('Density')
    plt.legend()
    plt.show()

    return g1_points, g2_points # Return the raw points


def plot_mir(mutation_rate_df, group_1_samples, group_2_samples, mir_name, p_value,
             group_1_label="Group1", group_2_label="Group2"):
    """
    Generates boxplots of mutation rates for a specific miRNA across groups.

    Parameters:
    -----------
    mutation_rate_df : pd.DataFrame
        DataFrame containing G>T mutation rates. Rows=miRNAs, Cols=Samples.
    group_1_samples : list
        List of sample IDs for group 1.
    group_2_samples : list
        List of sample IDs for group 2.
    mir_name : str
        Name of the miRNA to plot.
    p_value : float
        P-value from the differential mutation test to display in the title.
    """
    # Plotting the top hit
    data = pd.DataFrame({
        'Rate': pd.concat([mutation_rate_df.loc[mir_name, group_1_samples],
                           mutation_rate_df.loc[mir_name, group_2_samples]]),
        'Condition': [group_1_label]*len(group_1_samples) + [group_2_label]*len(group_2_samples)
    })
    plt.figure(figsize=(5, 6))
    sns.boxplot(x='Condition', y='Rate', data=data, palette='Set3')
    sns.stripplot(x='Condition', y='Rate', color='black', alpha=0.5, data=data)
    plt.title(f'Mutation Rate for {mir_name}\n(p-value: {p_value:.4f})')
    plt.ylabel('G>T / Total Reads')
    plt.show()
