import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import os
import pandas as pd
import importlib.resources

def plot_distances(distances_matrix, region, output_dir=None):
    # Ensure the input is a DataFrame to leverage labels
    if not isinstance(distances_matrix, pd.DataFrame):
        distances_matrix = pd.DataFrame(distances_matrix)

    # Load prototypes to identify which labels are user sequences
    try:
        user_input_cap = region.capitalize()
        if user_input_cap == 'Vp1':
            prototype_filename = 'vp1_test.csv'
        elif user_input_cap == 'Vp4/2':
            prototype_filename = 'prototypes.csv'
        else:
            prototype_filename = 'prototypes.csv'

        path_obj = importlib.resources.files('rhinotype.data').joinpath(prototype_filename)
        with importlib.resources.as_file(path_obj) as path:
            prototypes_df = pd.read_csv(path)
        proto_set = set(prototypes_df['Accession'])
    except Exception as e:
        print(f"Warning: Could not load prototype data for coloring distance labels: {e}")
        proto_set = set()

    # Plotting the heatmap
    num_sequences = distances_matrix.shape[0]

    # --- Dynamic Scaling for All Labels ---
    # Ensure every sequence label is shown.
    # Allocate fixed vertical/horizontal space per label to prevent overlap.
    scale_per_label = 0.25  # inches per label
    min_size = 15
    calc_size = num_sequences * scale_per_label
    fig_size = max(min_size, calc_size)
    # Cap size to prevent Matplotlib backend errors (limit is ~65k pixels)
    # 200 inches * 300 DPI = 60,000 pixels
    MAX_SIZE_INCHES = 200
    dpi_val = 300
    if fig_size > MAX_SIZE_INCHES:
        fig_size = MAX_SIZE_INCHES
        print(f"Warning: Heatmap size capped at {MAX_SIZE_INCHES} inches.")
        # Reduce font if compressed
        font_size = max(4, (MAX_SIZE_INCHES / num_sequences) * 50)
    else:
        font_size = 10
    plt.figure(figsize=(fig_size, fig_size * 0.9), dpi=dpi_val)
    # Force showing ALL labels with xticklabels=True, yticklabels=True
    ax = sns.heatmap(distances_matrix, cmap='YlOrRd', cbar=True, 
                xticklabels=True, yticklabels=True)
    plt.title(f"{region} genetic distances between sequences")

    # Rotate labels
    plt.xticks(rotation=90, fontsize=font_size)
    plt.yticks(rotation=0, fontsize=font_size)

    # --- Color User Sequence Labels Red ---
    # Apply to X axis
    for lbl in ax.get_xticklabels():
        if lbl.get_text() not in proto_set:
            lbl.set_color('red')
    # Apply to Y axis
    for lbl in ax.get_yticklabels():
        if lbl.get_text() not in proto_set:
            lbl.set_color('red')

    # Handle output directory
    if output_dir is None:
        output_dir = os.path.join(os.getcwd(), "figures")
    os.makedirs(output_dir, exist_ok=True)
    save_path = os.path.join(output_dir, "distances.png")
    # Use bbox_inches='tight' to prevent label truncation
    plt.savefig(save_path, bbox_inches='tight')
    # plt.show()
    print(f"Figure saved at: {save_path}")

if __name__ == "__main__":
    plot_distances()
