import pandas as pd
import numpy as np
import seaborn as sns
from turftopic import load_model
import matplotlib.pyplot as plt

def prepare_alignment_data(cfg):
    sample = pd.read_csv(f"{cfg.output_dir}/{cfg.output_file}")
    sample.columns = sample.columns.astype(str)

    model = load_model(cfg.data[cfg.level].topic_model)
    classification_cols = ['message_id', 'query', 'class']
    topic_cols = [str(c) for c in range(len(model.get_topics()))]

    print(classification_cols+topic_cols)
    print(cfg)
    return sample[classification_cols+topic_cols], model

def plot_heatmap(cfg):
    alignment_df, model = prepare_alignment_data(cfg)
    topic_cols = [str(c) for c in range(len(model.get_topics()))]
    classes = alignment_df['class'].unique()

    # Data matrix must have shape (len(y_edges)-1, len(x_edges)-1) -> (2, 3)
    data = np.array([(alignment_df[alignment_df['class']==cl][topic_cols]!=0).sum().values for cl in classes])
    col_labels = [f'topic_{k}' for k in range(len(model.get_topics()))]

    # 2. Configure the plot layout
    fig, ax = plt.subplots(figsize=(5, 3))
    sns.set_theme(font_scale=1.5) 
    # 3. Create the heatmap
    sns.heatmap(
        data,
        annot=True,  # Write data numbers inside each cell
        annot_kws={"size": 6},
        fmt=".2f",  # Format labels to 2 decimal places
        cmap="coolwarm",  # Use a blue-to-red color palette
        xticklabels=col_labels,  # Set X axis labels
        yticklabels=classes,  # Set Y axis labels
        linewidths=0.5,  # Add thin lines between cells
        square=True,  # Force cells to be perfect squares
    )

    # 2. Adjust axis titles
    # plt.xlabel("X Axis", fontsize=6)
    # plt.ylabel("Y Axis", fontsize=6)
    
    # # 3. Adjust axis tick labels (numbers on the grid edge)
    # ax.tick_params(labelsize=12)
    
    # 4. Adjust colorbar tick font size
    ax.figure.axes[-1].tick_params(labelsize=6)
    # 4. Show the visual output
    ax.set_title("Alignment Heatmap", fontsize=5, pad=6)
    ax.tick_params(axis='x', labelrotation=45, labelsize=6)
    ax.tick_params(axis='y', labelsize=6)
    
    return fig