data_engine.py

"""
Luke Cirne
Data Engine for ColabFold Wrapper
Functions for graphing and data analysis
Ma Lab
"""

import matplotlib.pyplot as plt
from matplotlib import ticker
import pandas as pd
import numpy as np
import random
import os
from collections import defaultdict


def compute_E(distances, R_0=51):
    for i in range(len(distances)):
        distances[i] = float(1 / (1 + (distances[i] / R_0)**6))
    return distances


def graph_output_accuracy(efficiencies: dict, bins=0.025, graph_name=None, N=None) -> str:
    # Collect and convert distances
    effs = np.array([float(d) for d in efficiencies.values()])
    total_strucs = len(effs)
    if not N:
        N = total_strucs

    # If bins is a float, treat it as bin width and generate edges
    if isinstance(bins, float) or isinstance(bins, int):
        min_d = effs.min()
        max_d = effs.max()
        bin_edges = np.arange(min_d, max_d + bins, bins)
        bin_edges = np.arange(0, 1 + bins, bins)
    else:
        # If bins is an array (from build_distribution), use it directly
        bin_edges = bins

    # Debug
    #print(effs.min(), effs.max())
    #print(bin_edges[:5], bin_edges[-5:])
    # Plot
    plt.figure(figsize=(8, 5))
    plt.hist(effs, bins=bin_edges, edgecolor="black", color="skyblue", label=f"Structures per Efficiency\ntotal structures: {total_strucs}")
    plt.title("CF Output Structures Separated by FRET Efficiency")
    plt.xlabel("FRET Efficiency")
    plt.ylabel("Frequency")
    plt.legend(title=f"N: {N}")

    plt.xticks(0, 1, 0.05)
    plt.tight_layout()

    # Gaussian curve (same x-range as histogram)
    y_exp = 0.291   # mean
    sigma = 0.083   # stdev
    x = np.linspace(bin_edges.min(), bin_edges.max(), 500)
    gaussian = (1 / (sigma * np.sqrt(2 * np.pi))) * np.exp(-0.5 * ((x - y_exp) / sigma) ** 2)
    # Scale to match histogram frequency
    gaussian_scaled = gaussian * len(effs) * (bin_edges[1] - bin_edges[0])
    plt.plot(x, gaussian_scaled, color="red", linewidth=2, label="Ideal Gaussian")

    # Save
    plot_name = "iteration_distances_hist"
    if graph_name:
        plot_name = graph_name
    plt.savefig(f"{plot_name}.png")
    return plot_name


def graph_output_accuracy_bar(efficiencies: dict, bins=0.0083, graph_name=None, N=None) -> str:
    """
    Plots a bar chart where each bar corresponds to a histogram bin.
    X values are bin centers, and Y values are counts of efficiencies in each bin.
    """

    # Convert dictionary values to numpy array
    filenames = np.array(list(efficiencies.keys()))
    effs = np.array([float(d) for d in efficiencies.values()])
    total_strucs = len(effs)
    if not N:
        N = total_strucs

    # Determine bin edges and centers
    if isinstance(bins, float) or isinstance(bins, int):
        min_d = effs.min()
        max_d = effs.max()
        bin_edges = np.arange(min_d, max_d + bins, bins)
        bin_centers = bin_edges[:-1] + bins / 2
    else:
        # If bins provided as array
        bin_edges = bins
        bin_centers = bin_edges[:-1] + (bin_edges[1] - bin_edges[0]) / 2

    print("===== GRAPH DEBUG =====")
    print("Number of efficiencies:", len(effs))
    print("Efficiencies:")
    print(effs)

    print("Bin edges:")
    print(bin_edges)

    # Count unique and duplicated files in each bin separately
    dupe_mask = np.char.find(np.char.lower(filenames.astype(str)), "dupe") >= 0
    unique_counts, _ = np.histogram(effs[~dupe_mask], bins=bin_edges)
    dupe_counts, _ = np.histogram(effs[dupe_mask], bins=bin_edges)

    print("Unique histogram counts:")
    print(unique_counts)
    print("Duplicate histogram counts:")
    print(dupe_counts)
    print("=======================")

    # --- Plot ---
    plt.figure(figsize=(8, 5))
    plt.bar(bin_centers, unique_counts, width=(bin_edges[1] - bin_edges[0]) * 0.9,
            color="mediumseagreen", edgecolor="black", label="Unique PDB files")
    plt.bar(bin_centers, dupe_counts, width=(bin_edges[1] - bin_edges[0]) * 0.9,
            bottom=unique_counts, color="coral", edgecolor="black", label="Duplicate PDB files")
    plt.title("CF Output Structures Separated by FRET Efficiency")
    plt.xlabel("FRET Efficiency")
    plt.ylabel("Frequency")
    plt.legend(title=f"N: {N}")

    # Set x-axis ticks
    #xticks = np.arange(bin_edges.min(), bin_edges.max() + 0.1, 0.1)
    xticks = np.arange(0, 1+0.1, 0.1)
    plt.xticks(xticks)
    plt.tight_layout()

    # --- Save ---
    plot_name = "iteration_distances_bar"
    if graph_name:
        plot_name = graph_name
    plt.savefig(f"{plot_name}.png")
    return plot_name


def build_distribution(
    file_eff_dict: dict,
    mean: float,
    std: float,
    bin_width: float = 0.0083,
    seed: int = None,
    n: int = None
) -> dict:

    """
    Selects file-efficiency pairs such that their histogram best follows a Gaussian distribution defined by the provided mean and standard deviation.

    Parameters:
        file_eff_dict : dict
            Dictionary of {filename: efficiency}, where efficiency is a float.
        mean : float
            Mean value for the target Gaussian distribution.
        std : float
            Standard deviation for the target Gaussian distribution.
        bin_width : float, optional
            Width of histogram bins. Default is 5.
        seed : int, optional
            Random seed for reproducibility. Default is 42.

    Returns
        dict: Dictionary of {filename: efficiency} containing the selected
            file-efficiency pairs adjusted to match the Gaussian distribution.
    """

    # Set seeds for reproducibility
    if seed is not None:
        np.random.seed(seed)
        random.seed(seed)
    # Extract efficiencies
    efficiencies = np.array(list(file_eff_dict.values()))
    filenames = np.array(list(file_eff_dict.keys()))

    # N represents the number of structures that will be returned
    N = n if n else len(efficiencies)

    # Define bin edges across observed range
    min_val, max_val = efficiencies.min(), efficiencies.max()
    #print(f"min_val: {min_val} max_val: {max_val}")
    #bins = np.arange(min_val, max_val + bin_width, bin_width)
    bins = np.arange(0, 1+bin_width, bin_width)

    # Bin assignments for each efficiency
    bin_indices = np.digitize(efficiencies, bins) - 1
    print(f"bin_indices: {bin_indices}")

    # Compute bin centers
    bin_centers = bins[:-1] + bin_width / 2

    # --- Compute Gaussian-based target counts ---
    gauss_probs = np.exp(-0.5 * ((bin_centers - mean) / std) ** 2)
    gauss_probs /= gauss_probs.sum()  # normalize
    target_counts = np.round(gauss_probs * N).astype(int)

    # Group files by bin
    bin_to_files = defaultdict(list)
    for fname, eff, bidx in zip(filenames, efficiencies, bin_indices):
        bin_to_files[bidx].append((fname, eff))

    # Collect selected filename-efficiency pairs
    selected = {}
    dupe_counter = defaultdict(int)

    mod_count = 0
    for bidx, desired_count in enumerate(target_counts):
        #print(f"desired_count: {desired_count}")
        available_files = bin_to_files.get(bidx, [])

        if desired_count == 0 or len(available_files) == 0:
            continue

        if len(available_files) >= desired_count:
            # Too many files, sample down
            chosen = random.sample(available_files, desired_count)
            mod_count += len(available_files) - desired_count
        else:
            # Too few files, duplicate as needed
            multiplier = -(-desired_count // len(available_files))  # ceiling division
            extended = available_files * multiplier
            chosen = random.sample(extended, desired_count)
            mod_count += desired_count - len(available_files)

        # Add chosen pairs with dupe suffixes if needed

        for fname, eff in chosen:
            if fname in selected:
                dupe_counter[fname] += 1
                name, ext = os.path.splitext(fname)
                new_fname = f"{name}_dupe{dupe_counter[fname]}{ext}"
                selected[new_fname] = eff
            else:
                selected[fname] = eff
        #print(len(chosen))
        #print(len(selected))

    return selected, bins, bin_centers, mod_count