import pandas as pd
import numpy as np
from scipy import stats
from scipy.stats import ks_2samp


def calculate_burden_richness(cpm_df, group, samples_df, threshold=1.0):
    """
    Calculates Total Burden (Sum) and Richness (Count > threshold).
        - Burden: The cumulative sum of CPM for all detected miRNAs per sample,
          representing the total magnitude of miRNA oxidation.
        - Richness: The count of detected miRNAs per sample above a specified threshold.
    output:
        - analysis_df: A dataframe with columns 'Burden', 'Richness', and 'group' for each sample.
        - stats_report: A dictionary containing the p-values and means
    parameters:
        - cpm_df: DataFrame with miRNAs as rows and samples as columns, containing CPM values.
        - group: The column name in samples_df that defines the groups to compare (e.g., 'Treatment', 'Activity').
        - samples_df: DataFrame containing sample metadata, including the 'Sample' column that matches cpm_df columns and the grouping column specified by 'group'.
        - threshold: The CPM threshold to consider a miRNA as "detected" for richness calculation (default is 1.0).
    """

    if group is None:
        raise ValueError("Group must be specified to calculate burden and richness.")
    if group not in samples_df.columns:
        raise ValueError(f"Group '{group}' not found in samples_df columns.")
    categories=[]
    unique_categories = samples_df[group].unique()
    for cat in unique_categories:
        categories.append((cat, samples_df['Sample'][samples_df[group] == cat].tolist()))
    results = []
    for group_name, sample_list in categories:
        # Calculate per-mouse metrics
        burden = cpm_df[sample_list].sum()
        richness = (cpm_df[sample_list] > threshold).sum()
        
        df = pd.DataFrame({
            'Burden': burden,
            'Richness': richness,
            'Group': group_name
        })
        results.append(df)
    
    analysis_df = pd.concat(results)
    
    # Statistics
    stats_report = {}
    # Instead of v, k, make it general so it is not Vehicle, KO, but any of categories
    for metric in ['Burden', 'Richness']:
        groups = []
        for cat in unique_categories:
            cat_values = analysis_df[analysis_df['Group'] == cat][metric]
            groups.append(cat_values)
        t_stat, p_val = stats.ttest_ind(*groups)
        stat_item = {'p_value': p_val, 't_stat': t_stat}
        for cat in unique_categories:
            cat_values = analysis_df[analysis_df['Group'] == cat][metric]
            stat_item[f'{cat}_mean'] = cat_values.mean()
        stats_report[metric] = stat_item
    return analysis_df, stats_report


def get_metadata_mapping(samples_df):
    """
    Standardizes the metadata, and creates a dictionary 
    pairing IP and NAT samples by individual.
    param samples_df: DataFrame containing sample metadata. Expected columns include:
    ['Sample', 'Individual', 'IP-Nat'] it can also have multpliple categorical columns.
    All columns that are not 'Sample', 'Individual', or 'IP-Nat' will be taken as 
    a categorical variable that can be used for subsetting the data later on.


    """
    df = samples_df.copy()
    categories = [col for col in df.columns if col not in ['Sample', 'Individual', 'IP-Nat']]


    # Create the Mouse Pairing Dictionary
    # We use 'Name on tube' to find the Mouse ID, removing ' IP' or ' Nat' suffixes
    pairs = {}
    for _, row in df.iterrows():
        sample_name = row['Sample']
        individual = row['Individual']
        ip_nat = row['IP-Nat']
        if individual not in pairs:
            pairs[individual] = {}
        pairs[individual][ip_nat.upper()] = sample_name
        for category in categories:
            if category not in pairs[individual]:
                pairs[individual][category] = row[category]

    return pairs


def calc_enrichment_matrix(cpm_df, mouse_pairs_dict, epsilon=0.1):
    """
    Pairs IP and NAT samples for each mouse to calculate propensity.
    The enrichment score is calculated as log2((IP + pseudocount) / (NAT + pseudocount)).
     - A pseudocount (epsilon) is added to both IP and NAT values to stabilize the ratio, especially for low CPM values.
     - The resulting enrichment scores are returned in a DataFrame with the same index as cpm_df and columns corresponding to each mouse ID.
     - This matrix can be used for downstream analyses, such as comparing enrichment across groups or performing statistical tests.
     - The function also includes error handling to skip any mice that do not have both IP and NAT samples present in the cpm_df.
     - The output DataFrame will have the same miRNA index as cpm_df and columns named after each mouse ID, containing the calculated enrichment scores.

     parameters:
     - cpm_df: DataFrame with miRNAs as rows and samples as columns, containing CPM values.
     - mouse_pairs_dict: Dictionary mapping each mouse ID to its corresponding IP and NAT sample names, as well as any additional categorical metadata.
     - epsilon: A small pseudocount added to both IP and NAT values to prevent division by zero and stabilize the log ratio (default is 0.1).
     returns:
     - enrichment_df: DataFrame with miRNAs as rows and mouse IDs as columns, containing the calculated enrichment scores (log2 ratio of IP to NAT).
    """
    enrichment_scores = {}
    
    for mouse_id, samples in mouse_pairs_dict.items():
        ip_col = samples['IP']
        nat_col = samples['NAT']
        
        # Calculate log2 ratio with pseudocount to stabilize low values
        if ip_col not in cpm_df.columns or nat_col not in cpm_df.columns:
            print(f"Warning: Missing columns for {mouse_id}. Skipping.")
            continue
        ip_vals = cpm_df[ip_col] + epsilon
        nat_vals = cpm_df[nat_col] + epsilon

        enrichment_scores[mouse_id] = np.log2(ip_vals / nat_vals +1)
        
    return pd.DataFrame(enrichment_scores, index=cpm_df.index)


def get_enrichment_score(ip_counts, nat_counts):
    """Calculate enrichment score as log2((IP + pseudocount) / (NAT + pseudocount)) as Gibran used
    params:
    - ip_counts: Series or array of CPM values for the IP sample.
    - nat_counts: Series or array of CPM values for the NAT sample.
     returns:
    - enrichment: Series or array of enrichment scores calculated as log2((IP + pseudocount) / (NAT + pseudocount)).
     Note: A pseudocount is added to both IP and NAT counts to prevent division by zero and stabilize the log ratio, especially for low CPM values.
     The resulting enrichment scores can be used for downstream analyses, such as comparing enrichment across groups or performing statistical tests.
     The function assumes that ip_counts and nat_counts are aligned and that they contain non-negative values representing CPM.
    """
    # Add a pseudocount to avoid division by zero
    pseudocount = 0.1
    enrichment = (ip_counts + pseudocount) / (nat_counts + pseudocount)
    enrichment = np.log2(enrichment)
    return enrichment


def filter_enrichment_noise(enrichment_df, cpm_df, mouse_pairs_dict, threshold=5.0):
    """
    Removes miRNAs where the average signal (IP + NAT) is below a threshold.
    Crucial for preventing zero-inflation from hiding real medians.
    """
    all_ip_cols = [v['IP'] for v in mouse_pairs_dict.values()]
    all_nat_cols = [v['NAT'] for v in mouse_pairs_dict.values()]
    
    # Calculate mean combined signal
    avg_signal = (cpm_df[all_ip_cols].mean(axis=1) + cpm_df[all_nat_cols].mean(axis=1))
    mask = avg_signal > threshold
    
    filtered_df = enrichment_df[mask].copy()
    print(f"Filtering: {len(filtered_df)} miRNAs retained with signal > {threshold} CPM.")
    return filtered_df


def run_distribution_tests(enrichment_df, group_1_individuals, group_2_individuals):
    """
    Performs dual-level statistics: Per-mouse averages and Pooled KS-test.
    """

    # remove any samples that are not in the enrichment_df columns (in case of missing pairs)
    group_1_individuals = [m for m in group_1_individuals if m in enrichment_df.columns]
    group_2_individuals = [m for m in group_2_individuals if m in enrichment_df.columns]
    # 1. Level 1: Per-mouse consistency (T-test)
    group_1_means = enrichment_df[group_1_individuals].mean()
    group_2_means = enrichment_df[group_2_individuals].mean()
    t_stat, p_ttest = stats.ttest_ind(group_1_means, group_2_means)
    
    # 2. Level 2: Transcriptome-wide shift (KS-test)
    # Melt data to pool every miRNA point per genotype
    group_1_points = enrichment_df[group_1_individuals].values.flatten()
    group_2_points = enrichment_df[group_2_individuals].values.flatten()
    
    ks_stat, p_ks = ks_2samp(group_1_points, group_2_points)
    
    stats_report = {
        'ttest_p': p_ttest,
        'ks_p': p_ks,
        'ks_stat': ks_stat,
        'group_1_median': np.median(group_1_points),
        'group_2_median': np.median(group_2_points)
    }
    
    return stats_report


def calculate_gt_rates(mismatch_df, sample_cols):
    """
    Normalizes G>T counts by Total Reads to get a 'Damage Rate'.
    """
    # 1. Numerator: Sum of G>T counts per miRNA
    gt_only = mismatch_df[mismatch_df['pos:mut'].str.contains('G>T', na=False)]
    numerator = gt_only.groupby('miRNA name')[sample_cols].sum()
    
    # 2. Denominator: Sum of all reads (PM + all Mismatches) per miRNA
    denominator = mismatch_df.groupby('miRNA name')[sample_cols].sum()
    
    # 3. Calculate Rate (add 1 to denominator to avoid div by zero)
    rate_df = numerator / (denominator + 1)
    return rate_df.fillna(0)


def extract_gt_positions(mismatch_df):
    """
    Parses 'pos:mut' column (e.g., '5:G>T,12:A>C') to find the first G>T position.
    """
    def get_first_gt(val):
        if pd.isna(val) or 'G>T' not in str(val): return np.nan
        # Splits by comma for multiple mutations, takes first G>T
        for part in str(val).split(','):
            if 'G>T' in part:
                # Extract digits only from the part before ':'
                pos = ''.join(filter(str.isdigit, part.split(':')[0]))
                return int(pos) if pos else np.nan
        return np.nan

    df = mismatch_df.copy()
    df['gt_position'] = df['pos:mut'].apply(get_first_gt)
    return df


def run_diff_mutation_test(rate_df, group_1_samples, group_2_samples, min_reads_df=None, min_threshold=50):
    """
    Compares mutation rates per miRNA using Welch's T-test.
    """
    results = []
    # Optional: Filter for miRNAs that have enough total reads to be reliable
    if min_reads_df is not None:
        valid_mirs = min_reads_df[list(group_1_samples) + list(group_2_samples)].sum(axis=1) > min_threshold
        rate_df = rate_df[rate_df.index.isin(valid_mirs[valid_mirs].index)]

    for mirna, row in rate_df.iterrows():
        g1_vals = row[group_1_samples].values
        g2_vals = row[group_2_samples].values
        
        t_stat, p_val = stats.ttest_ind(g1_vals, g2_vals, equal_var=False)
        
        results.append({
            'miRNA': mirna,
            'Rate_G1': g1_vals.mean(),
            'Rate_G2': g2_vals.mean(),
            'Delta': g2_vals.mean() - g1_vals.mean(),
            'p_value': p_val
        })
        
    return pd.DataFrame(results).sort_values('p_value').set_index('miRNA')


def run_mutation_rate_burden(rate_df, group_1_samples, group_2_samples, 
                             group_1_label="G1", group_2_label="G2",
                               condition="Exercise", title="Global G>T Mutation Burden"):
    """
    Compares the global G>T mutation burden between two groups of samples.
    The mutation burden is calculated as the mean G>T mutation rate across all miRNAs 
    for each sample, normalized by total reads.

    Parameters:
    -----------
    rate_df : pd.DataFrame
        DataFrame containing G>T mutation rates for each miRNA (rows) and sample (columns
        The values should be the G>T mutation rate (e.g., G>T counts / total reads).
    group_1_samples : list with the name of the SAMPLES (not individuals) of the first group (e.g., Vehicle)
    group_2_samples : list with the name of the SAMPLES (not individuals) of the second group (e.g., KO)
    group_1_label : str
        Label for the first group (default: "G1").
    group_2_label : str
        Label for the second group (default: "G2").
    condition : str
        Description of the condition being compared (used for plot titles, default: "Exercise").
    title: str
        Title for the plot (default: "Global G>T Mutation Burden").


    returns:
        eval_dict : dict
            Dictionary containing the t-statistic, p-value, group means, standard deviations, and counts.
    """
    import matplotlib.pyplot as plt
    import seaborn as sns

    valid_samples = list(group_1_samples) + list(group_2_samples)
    if not all(sample in rate_df.columns for sample in valid_samples):
        missing = [sample for sample in valid_samples if sample not in rate_df.columns]
        raise ValueError(f"The following samples are missing from the rate DataFrame: {missing}")

    # We take the mean mutation rate across all miRNAs for each mouse
    mouse_burden = rate_df[valid_samples].mean().to_frame('G2T_Mutation_Rate')
    mouse_burden[condition] = [group_1_label if col in group_1_samples else group_2_label for col in mouse_burden.index]
    
    # Stats
    v_vals = mouse_burden[mouse_burden[condition] == group_1_label]['G2T_Mutation_Rate']
    k_vals = mouse_burden[mouse_burden[condition] == group_2_label]['G2T_Mutation_Rate']
    t_stat, p_val = stats.ttest_ind(v_vals, k_vals)
    
    print(f"--- G>T Mutation Rate Burden ({condition}) ---")
    print(mouse_burden.groupby(condition)['G2T_Mutation_Rate'].agg(['mean', 'std', 'count']))
    print(f"\nT-test p-value: {p_val:.4f}")
    print(f"\nT-test t-stat: {t_stat:.4f}")

    # Plot
    plt.figure(figsize=(6, 6))
    sns.boxplot(x=condition, y='G2T_Mutation_Rate', data=mouse_burden, palette='Set2')
    sns.stripplot(x=condition, y='G2T_Mutation_Rate', data=mouse_burden, color='black', alpha=0.5)
    plt.title(f'{title}\n({condition})')
    plt.ylabel('Mutation Rate (G>T / Total Reads +1 )')
    plt.show()

    eval_dict = {
        't_statistic': t_stat,
        'p_value': p_val,
        'group_1_mean': v_vals.mean(),
        'group_2_mean': k_vals.mean(),
        'group_1_std': v_vals.std(),
        'group_2_std': k_vals.std(),
        'group_1_count': v_vals.count(),
        'group_2_count': k_vals.count()
    }
    return eval_dict


def find_differentially_oxidized_mirs(rate_df, total_counts_df, group_1_samples, group_2_samples, min_total_reads=50):
    """
    Identifies miRNAs that are differentially oxidized between two groups of samples based on G>T mutation rates.

    Parameters:
    -----------
    rate_df : pd.DataFrame
        DataFrame containing G>T mutation rates for each miRNA (rows) and sample (columns).
    total_counts_df : pd.DataFrame
        DataFrame containing total read counts for each miRNA (rows) and sample (columns).
    group_1_samples : list
        List of sample names for the first group (e.g., Vehicle).
    group_2_samples : list
        List of sample names for the second group (e.g., KO).
    min_total_reads : int
        Minimum total reads required for a miRNA to be considered (default: 50).

    Returns:
    --------
    pd.DataFrame
        DataFrame containing the results of the differential analysis, including miRNA names, mutation rates, and p-values.
    """
    results = []
    
    # Filter for miRNAs that have enough "material" to be statistically valid
    # We look at the sum of total reads across the relevant samples
    mask = total_counts_df[list(group_1_samples) + list(group_2_samples)].sum(axis=1) > min_total_reads
    test_df = rate_df[mask]
    
    print(f"Testing {len(test_df)} miRNAs that passed the {min_total_reads} total reads threshold.")

    for mirna, row in test_df.iterrows():
        g1_vals = row[group_1_samples].values
        g2_vals = row[group_2_samples].values
        
        # Calculate Means
        mean_g1 = g1_vals.mean()
        mean_g2 = g2_vals.mean()
        
        # Calculate Change (Delta)
        delta = mean_g2 - mean_g1
        
        # Statistical Test
        # We use nan_policy='omit' just in case
        t_stat, p_val = stats.ttest_ind(g1_vals, g2_vals, equal_var=False)
        print(f"Results for {mirna}: t-statistic = {t_stat}, p-value = {p_val}")
        results.append({
            'miRNA': mirna,
            'Rate_G1': mean_g1,
            'Rate_G2': mean_g2,
            'Delta_Rate': delta,
            'p_value': p_val
        })
        print(len(results))
    print("Differentially oxidized miRNAs found:", len(results))
    print(results)
    res_df = pd.DataFrame(results)

    
    # Simple Multiple Testing Correction
    #res_df['padj'] = res_df['p_value'] * len(res_df) # Basic Bonferroni

    
    return res_df.sort_values('p_value')


def extract_first_position(oxi_mir_df):
    """
    Extracts the first position of G>T mutations from the 'pos:mut' column in the input DataFrame.

    Parameters:
    -----------
    oxi_mir_df : pd.DataFrame
        DataFrame containing a 'pos:mut' column with mutation position information.

    Returns:
    --------
    list
        A list of integers representing the first positions of G>T mutations for each miRNA.
    """
    if 'pos:mut' not in oxi_mir_df.columns:
        raise ValueError("The 'pos:mut' column is missing from the DataFrame.")
    first_gt = []
    for pos_mut in oxi_mir_df['pos:mut']:
        if 'G>T' in pos_mut:
            # Extract the first position number before "G>T"
            pos = pos_mut.split('G>T')[0]
            pos = pos.split(',')[-1]  # Get the part after the last ','
            # Remove any non-digit characters
            pos = ''.join(filter(str.isdigit, pos))
            first_gt.append(int(pos) if pos else np.nan)
        else:
            first_gt.append(np.nan)
    return first_gt
