import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker

from src.utils.vis import figure, subplots

def save_group_lists_fasta(loci_group, cap_group, data_save_path):
    loci_file = os.path.join(data_save_path, 'loci_group_list.txt')
    with open(loci_file, 'w') as f:
        for group in sorted(loci_group['group'].unique()):
            f.write(f">loci_group_{group}\n")
            group_loci = loci_group[loci_group['group'] == group]['loci'].values
            if len(group_loci) > 0:
                # Extract unique prefixes before '-' for loci only
                gene_names = extract_unique_genes(group_loci)
                for gene in gene_names:
                    f.write(f"{gene}\n")
            else:
                f.write("# Empty group\n")
            f.write("\n")
    
    cap_file = os.path.join(data_save_path, 'cap_group_list.txt')
    with open(cap_file, 'w') as f:
        for group in sorted(cap_group['group'].unique()):
            f.write(f">cap_group_{group}\n")
            group_caps = cap_group[cap_group['group'] == group]['cap'].values
            if len(group_caps) > 0:
                for cap in group_caps:
                    f.write(f"{cap}\n")
            else:
                f.write("# Empty group\n")
            f.write("\n")

def extract_unique_genes(loci_list):
    # Extract gene names before "-"
    unique_genes = set()
    for loci in loci_list:
        if '-' in loci:
            unique_genes.add(loci.split('-')[0])
        else:
            unique_genes.add(loci)
    return unique_genes

def iterative_refinement(diff_mat, loci_assignment, cap_assignment, loci_list = None, cap_name = None, top_group_percent = 0.3, top_item_percent = 0.3, n_iter = 50):
    loci_list = pd.Series(loci_list)
    cap_name = pd.Series(cap_name)
    loci_assignment = pd.Series(loci_assignment)
    cap_assignment = pd.Series(cap_assignment)

    loci_rank_df = pd.DataFrame({'loci': loci_list, 'value': np.mean(diff_mat, axis=1)})
    sorted_loci_rank_df = loci_rank_df.sort_values(by='value', ascending=False)
    sorted_loci_rank_df['rank'] = np.arange(len(sorted_loci_rank_df))
    loci_rank_df = sorted_loci_rank_df.sort_index()
    loci_rank = loci_rank_df['rank']

    cap_rank_df = pd.DataFrame({'cap': cap_name, 'value': np.mean(diff_mat, axis=0)})
    sorted_cap_rank_df = cap_rank_df.sort_values(by='value', ascending=False)
    sorted_cap_rank_df['rank'] = np.arange(len(sorted_cap_rank_df))
    cap_rank_df = sorted_cap_rank_df.sort_index()
    cap_rank = cap_rank_df['rank']

    for i in range(n_iter):
        # Refine loci assignment
        loci_rank_new = refine_rank(diff_mat, cap_assignment, cap_rank, top_group_percent, top_item_percent)
        loci_rank_diff = np.sum(np.abs(loci_rank_new - loci_rank))
        loci_rank = loci_rank_new
        # Refine cap assignment
        cap_rank_new = refine_rank(diff_mat.T, loci_assignment, loci_rank, top_group_percent, top_item_percent)
        cap_rank_diff = np.sum(np.abs(cap_rank_new - cap_rank))
        cap_rank = cap_rank_new
        print(f'Loci rank difference: {loci_rank_diff} \n cap rank difference: {cap_rank_diff}')
        if loci_rank_diff + cap_rank_diff < 1e-6:
            break
    return loci_rank, cap_rank

def refine_rank(diff_mat, other_assignment, other_rank, top_group_percent, top_item_percent):
    total_groups = len(np.unique(other_assignment))
    # Get the top percent of data points from other_rank
    top_group_idx = int(total_groups * top_group_percent)
    # Select the top group_idx groups
    top_groups = other_assignment[other_assignment <= top_group_idx] # Keep to top groups in other_assignment
    top_group_rank = other_rank[top_groups.index] # Get the rank of the top groups

    top_other_rank = other_rank.sort_values(ascending=True)[:int(len(other_rank) * top_item_percent)] # Get the top item_percent of data points in top groups
    top_other_rank_idx = top_other_rank.index

    # Get the data points from diff_mat for ranking the current rank
    top_other_rank_data = diff_mat[:, top_other_rank_idx]
    # Get the mean of the top other rank data
    top_other_rank_data_mean = np.mean(top_other_rank_data, axis=1)
    # Get current rank
    current_rank_df = pd.DataFrame({'value': top_other_rank_data_mean})
    current_rank_df = current_rank_df.sort_values(by='value', ascending=False)
    current_rank_df['rank'] = np.arange(len(current_rank_df))
    current_rank = current_rank_df['rank'].sort_index()
    return current_rank


def plot_signature_heatmap_full(plot_dict, block_size = 400, vmin = -10, vmax = 10, show_ticks = True):
    plot_mat = plot_dict['plot_mat']
    loci_k = plot_dict['loci_k']
    cap_k = plot_dict['cap_k']
    loci_group = plot_dict['loci_group']
    cap_group = plot_dict['cap_group']
    data_save_path = plot_dict['data_save_path']
    x_name = plot_dict['x_name']
    y_name = plot_dict['y_name']
    x_annotation = plot_dict['x_annotation']
    y_annotation = plot_dict['y_annotation']
    title_str = plot_dict['title_str']
    direction = plot_dict['direction']

    plot_mat, loci_k, cap_k, loci_group, cap_group, x_name, y_name, x_annotation, y_annotation = invert_heatmap(plot_mat, loci_k, cap_k, loci_group, cap_group, x_name, y_name, x_annotation, y_annotation)

    save_name = f'{data_save_path}/signature_heatmap_full.pdf'
    #width = len(cap_group) * 1
    #height = len(loci_group) * 1
    width, height = 45, 45
    fig, ax = subplots(figsize=(width, height))
    
    #ax.imshow(plot_mat, cmap='coolwarm', aspect='auto', vmin=-10, vmax=10, interpolation='None', rasterized=False)
    # Instead use pcolormesh to plot heatmap
    if direction == 'upregulated':
        plt.gca().invert_yaxis() # This is to invert the y-axis for down heatmap
    #ax.pcolormesh(plot_mat, cmap='coolwarm', vmin=-10, vmax=10)
    #ax.pcolormesh(plot_mat, cmap='coolwarm', vmin=-10, vmax=10)
    # Draw lines to separate groups
    from skimage.measure import block_reduce
    width, height = plot_mat.shape
    x_block_size = max(width // block_size, 1)
    y_block_size = max(height // block_size, 1)
    heatmap = block_reduce(plot_mat, (x_block_size, y_block_size), np.max)

    yedges = np.linspace(0, 1, heatmap.shape[0])
    xedges = np.linspace(0, 1, heatmap.shape[1])
    #ax.imshow(heatmap, cmap='viridis', aspect='auto', interpolation='none', rasterized=False, vmin=0, vmax=2)
    if vmin is None:
        vmin = np.min(heatmap)
    if vmax is None:
        vmax = np.max(heatmap)
    ax.pcolormesh(xedges, yedges, heatmap, cmap='RdBu_r', rasterized=False, vmin=vmin, vmax=vmax)
    #ax.imshow(heatmap, cmap='RdBu_r', aspect='auto', interpolation='none', rasterized=False, vmin=-10, vmax=10)
    width_block = width // block_size
    height_block = height // block_size

    y_delta = yedges[1] - yedges[0]
    x_delta = xedges[1] - xedges[0]

    for i in range(0, loci_k):
        group_indices = np.where(loci_group['group'] == i)[0]
        if len(group_indices) > 0:  # Check if group exists and is not empty
            h_idx = group_indices[-1] / width
            # Select the closest yedge
            h_idx = yedges[np.argmin(np.abs(yedges - h_idx))]
            h_idx = h_idx + y_delta / 2
            ax.axhline(h_idx, color='black', linewidth=0.2, linestyle='--', alpha = 0.5)
    for i in range(0, cap_k):
        group_indices = np.where(cap_group['group'] == i)[0]
        if len(group_indices) > 0:  # Check if group exists and is not empty
            v_idx = group_indices[-1] / height
            # Select the closest xedge
            v_idx = np.argmin(np.abs(xedges - v_idx))
            v_idx = xedges[v_idx] - x_delta / 2
            ax.axvline(v_idx, color='black', linewidth=0.2, linestyle='--', alpha = 0.5)
    '''
    ax.set_xlabel('Position')
    ax.set_ylabel('CAP')
    '''
    if show_ticks:
        total_ticks = 10
        '''
        ax.set_yticks(yedges[::len(yedges) // total_ticks][:total_ticks])
        ax.set_xticks(xedges[::len(xedges) // total_ticks][:total_ticks])
        ax.set_yticklabels(range(width)[::width // total_ticks][:total_ticks])
        ax.set_xticklabels(range(height)[::height // total_ticks][:total_ticks])
        # Rotate x tick labels
        ax.set_xticklabels(ax.get_xticklabels(), rotation=90)
        '''
        ax.xaxis.set_major_locator(mticker.MaxNLocator(nbins=total_ticks))
        ax.yaxis.set_major_locator(mticker.MaxNLocator(nbins=total_ticks)) ##hua added
        ax.set_xticklabels(ax.get_xticklabels(), rotation=90)
    else:
        ax.set_xticks([])
        ax.set_yticks([])
    # remove x spine
    ax.spines['top'].set_visible(True)
    ax.spines['right'].set_visible(True)
    ax.set_xlabel(x_name)
    ax.set_ylabel(y_name)
    #ax.set_xticks(np.arange(len(cap_group)) + 0.5)
    #ax.set_yticks(np.arange(len(loci_group)) + 0.5)
    #ax.set_xticklabels(x_annotation, rotation=90)
    #ax.set_yticklabels(y_annotation)
    # Change label font size
    #ax.xaxis.set_tick_params(labelsize=2)
    #ax.yaxis.set_tick_params(labelsize=2)
    ax.set_title(title_str)
    plt.savefig(save_name, bbox_inches='tight')
    plt.close()

def plot_signature_heatmap_group(plot_dict):
    plot_mat = plot_dict['plot_mat']
    loci_k = plot_dict['loci_k']
    cap_k = plot_dict['cap_k']
    loci_group = plot_dict['loci_group']
    cap_group = plot_dict['cap_group']
    data_save_path = plot_dict['data_save_path']
    x_name = plot_dict['x_name']
    y_name = plot_dict['y_name']
    x_annotation = plot_dict['x_annotation']    
    y_annotation = plot_dict['y_annotation']
    title_str = plot_dict['title_str']
    
    # Plot heatmap of groups
    group_mean_mat = np.zeros((loci_k, cap_k))
    for i in range(loci_k):
        for j in range(cap_k):
            loci_mask = loci_group['group'] == i
            cap_mask = cap_group['group'] == j
            
            # Check if both groups have members
            if np.any(loci_mask) and np.any(cap_mask):
                group_mean_mat[i, j] = np.mean(plot_mat[loci_mask, :][:, cap_mask])
            else:
                group_mean_mat[i, j] = 0  # Set to 0 if either group is empty
    group_mean_df = pd.DataFrame(
        group_mean_mat,
        index = [f'loci_group_{i}' for i in range(loci_k)], 
        columns = [f'cap_group_{i}' for i in range(cap_k)]
        )
    group_mean_df.to_csv(f'{data_save_path}/group_mean_matrix.csv')
    save_name = f'{data_save_path}/signature_heatmap_group.pdf'

    plot_mat, loci_k, cap_k, loci_group, cap_group, x_name, y_name, x_annotation, y_annotation = invert_heatmap(group_mean_mat, loci_k, cap_k, loci_group, cap_group, x_name, y_name, x_annotation, y_annotation)
    fig, ax = subplots(figsize=(70, 70))
    ax.imshow(plot_mat, cmap='coolwarm', aspect='auto', vmin=-5, vmax=5, interpolation=None)
    ax.set_xlabel(x_name)
    ax.set_ylabel(y_name)
    ax.set_xticks(np.arange(cap_k))
    ax.set_yticks(np.arange(loci_k))
    ax.set_xticklabels(np.arange(cap_k))
    ax.set_yticklabels(np.arange(loci_k))
    ax.set_title(title_str)
    plt.savefig(save_name)
    plt.close()

def invert_heatmap(plot_mat, loci_k, cap_k, loci_group, cap_group, x_name, y_name, x_annotation, y_annotation):
    return plot_mat.T, cap_k, loci_k, cap_group, loci_group, y_name, x_name, y_annotation, x_annotation

def plot_group_sizes(loci_group, cap_group, data_save_path):
    # Get all unique groups from both loci and cap groups
    all_groups = sorted(set(loci_group['group'].unique()) | set(cap_group['group'].unique()))
    
    # Create size dataframes with all groups, filling missing with 0
    loci_sizes = loci_group.groupby('group').size().reset_index(name='size')
    cap_sizes = cap_group.groupby('group').size().reset_index(name='size')
    
    # Ensure all groups are represented, filling missing groups with size 0
    loci_sizes = loci_sizes.set_index('group').reindex(all_groups, fill_value=0).reset_index()
    cap_sizes = cap_sizes.set_index('group').reindex(all_groups, fill_value=0).reset_index()
    
    # Get unique genes in each loci group
    gene_sizes = []
    for group in all_groups:
        # Loci unique prefixes
        group_loci = loci_group[loci_group['group'] == group]['loci'].values
        if len(group_loci) > 0:
            unique_genes = extract_unique_genes(group_loci)
            gene_sizes.append(len(unique_genes))
        else:
            gene_sizes.append(0)  # Empty group

    plt.figure(figsize=(10, 6))
    bar_width = 0.25
    x = np.arange(len(all_groups))
    plt.bar(x - bar_width, loci_sizes['size'], width=bar_width, label='Loci', color='#5EB1BF')
    plt.bar(x, gene_sizes, width=bar_width, label='Unique genes', color='#042A2B')
    plt.bar(x + bar_width, cap_sizes['size'], width=bar_width, label='CAPs', color='#EF7B45')
    plt.xlabel('Group')
    plt.ylabel('Size')
    plt.xticks(x, all_groups)
    plt.title('Group Sizes')
    plt.legend()
    plt.savefig(os.path.join(data_save_path, 'group_sizes.png'))
    plt.close()

def compute_jaccard_similarity(set1, set2):
    if len(set1 | set2) == 0:
        return 0.0
    return len(set1 & set2) / len(set1 | set2)

def compute_jaccard_matrix(group_df1, group_df2, entity_col, group_col):
    groups_1 = sorted(group_df1[group_col].unique())
    groups_2 = sorted(group_df2[group_col].unique())
    
    similarity_matrix = np.zeros((len(groups_1), len(groups_2)))
    
    for i, g1 in enumerate(groups_1):
        for j, g2 in enumerate(groups_2):
            set_1 = set(group_df1[group_df1[group_col] == g1][entity_col])
            set_2 = set(group_df2[group_df2[group_col] == g2][entity_col])
            similarity_matrix[i, j] = compute_jaccard_similarity(set_1, set_2)
    
    return similarity_matrix, groups_1, groups_2