import numpy as np
import pandas as pd
from scipy.cluster.hierarchy import dendrogram, linkage
from scipy.spatial.distance import squareform
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import os
import importlib.resources
import random

def plot_tree(distance_matrix, region, output_dir=None):
    # Load prototypes to identify species/genotypes and distinguish 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)
        # Create mappings
        # Accession -> Species (A, B, C)
        proto_species = dict(zip(prototypes_df['Accession'], prototypes_df['Species']))
        # Accession -> Genotype (e.g. A1, B101)
        if 'Genotype' in prototypes_df.columns:
            proto_genotype = dict(zip(prototypes_df['Accession'], prototypes_df['Genotype']))
        else:
            proto_genotype = {}
        proto_set = set(prototypes_df['Accession'])
        # Generate Genotype Colors
        # Get unique genotypes sorted, filtering out NaNs
        raw_genotypes = list(set(proto_genotype.values()))
        unique_genotypes = sorted([g for g in raw_genotypes if pd.notna(g) and str(g).strip() != ''])
        num_genotypes = len(unique_genotypes)
        indices = list(np.linspace(0, 1, num_genotypes))
        random.seed(42) 
        random.shuffle(indices)
        cmap = cm.get_cmap('nipy_spectral')
        genotype_colors = {
            g: cm.colors.to_hex(cmap(idx)) 
            for g, idx in zip(unique_genotypes, indices)
        }

    except Exception as e:
        print(f"Warning: Could not load prototype data for coloring: {e}")
        proto_species = {}
        proto_genotype = {}
        proto_set = set()
        genotype_colors = {}
    # Extract labels if the input is a pandas DataFrame
    labels = None
    if hasattr(distance_matrix, 'index'):
        labels = distance_matrix.index.tolist()
    else:
        labels = [str(i) for i in range(len(distance_matrix))]
    # Identify user sequences and assign Species AND Genotype to leaves
    leaves_species = [] 
    leaves_genotype = []
    for i, label in enumerate(labels):
        if label in proto_set:
            # Prototype
            s = proto_species.get(label)
            g = proto_genotype.get(label)
            # Clean species/genotype
            s = str(s).strip() if pd.notna(s) else None
            g = str(g).strip() if pd.notna(g) else None
            leaves_species.append({s} if s else set())
            leaves_genotype.append({g} if g else set())
        else:
            # User sequence - treat as neutral (wildcard)
            leaves_species.append(set())
            leaves_genotype.append(set())
    # Convert the data to a numpy array for calculations
    dm_array = np.array(distance_matrix)
    condensed_distance_matrix = squareform(dm_array)
    hc = linkage(condensed_distance_matrix, method='complete')

    # Determine Clade Colors (Species & Genotype)
    cluster_species = leaves_species[:]
    cluster_genotype = leaves_genotype[:]
    # Precompute sets for all clusters
    for i in range(len(hc)):
        c1_idx = int(hc[i, 0])
        c2_idx = int(hc[i, 1])
        # Union of properties present in children
        union_s = cluster_species[c1_idx] | cluster_species[c2_idx]
        union_g = cluster_genotype[c1_idx] | cluster_genotype[c2_idx]
        cluster_species.append(union_s)
        cluster_genotype.append(union_g)
    # Define color function
    def get_color(k):
        # 1. Check for Specific Genotype Uniformity
        g_set = cluster_genotype[k]
        # Filter out invalid genotypes from the set check
        valid_g = {g for g in g_set if g and pd.notna(g)}
        if len(valid_g) == 1:
            g_name = list(valid_g)[0]
            if g_name in genotype_colors:
                return genotype_colors[g_name]
        # 2. Fallback to Species Uniformity (Major Clades)
        s_set = cluster_species[k]
        # Filter out invalid species
        valid_s = {s for s in s_set if s and pd.notna(s)}
        if valid_s == {'A'}:
            return 'blue'
        elif valid_s == {'B'}:
            return 'red'
        elif valid_s == {'C'}:
            return 'green'
        else:
            return 'black'
    # Plot the dendrogram
    num_leaves = len(dm_array)
    scale_per_leaf = 0.2 if num_leaves < 200 else 0.1
    fig_height = max(8, num_leaves * scale_per_leaf)

    MAX_HEIGHT_INCHES = 200
    if fig_height > MAX_HEIGHT_INCHES:
        fig_height = MAX_HEIGHT_INCHES
        print(f"Warning: Tree height capped at {MAX_HEIGHT_INCHES} inches.")
    dpi = 600
    if num_leaves > 200: dpi = 300
    if num_leaves > 1000: dpi = 150
    plt.figure(figsize=(15, fig_height), dpi=dpi)
    ddata = dendrogram(hc, 
                       orientation='left', 
                       labels=labels, 
                       leaf_font_size=8,
                       link_color_func=get_color)
    plt.title(f"{region} simple tree ({num_leaves} sequences)")
    plt.xlabel("Genetic distance")
    plt.ylabel("")
    # Color Leaf Labels
    # Highlight user sequences in RED
    ax = plt.gca()
    tick_labels = ax.get_yticklabels()
    for lbl in tick_labels:
        text = lbl.get_text()
        if text not in proto_set:
            # User sequence -> Red
            lbl.set_color('red')

    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, "tree.png")

    plt.savefig(save_path, bbox_inches='tight')
    # plt.show()

    print(f"Figure saved at: {save_path}")

if __name__ == "__main__":
    plot_tree()
