from lookup_sparsity_blob import SparseBlob
import matplotlib.pyplot as plt
from glob import glob

def plot_layer_hist(df, row_id, ax, type=None):
    if type is None:
        return "pls specify type='widest' or 'narrowest'"
    
    cols = ['layer_id', 'shape', 'sparsity','tile_limit', 'tile_min', 'tile_med', 'tile_max', 'tile_shape', 'n_tile']
    meta = df.loc[row_id, cols]
    tensor = sb._blob[meta.layer_id]['tile_sparsity']

    l = meta.layer_id
    l_shape = str(tuple(meta['shape']))
    l_sparsity = meta.sparsity*100
    hist_limit = f"min: {meta.tile_min:.2f}, median: {meta.tile_med:.2f}, max: {meta.tile_max:.2f}"
    l_cfg = f"{meta.n_tile} tiles of {meta.tile_shape}"

    ax.hist(tensor.flatten().cpu().numpy(), bins=25, edgecolor='white', color='blue')
    ax.set_xlabel('Tile Sparsity')
    ax.set_ylabel('Frequency')
    ax.set_title(f'{l} {l_shape}\nlayer_sparsity: {l_sparsity:2.0f}%\n\ntile_sparsity ({type}):\n{hist_limit}\n\n{l_cfg}')

    ax.grid(True, axis='y', linestyle='--', alpha=0.3)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.tick_params(axis='both', which='major', labelsize=12)


sparse_model_list = sorted(glob("blob.sparsity*"))

for model_blob in sparse_model_list:
    plot_name = model_blob.replace("blob.sparsity.", "tile_sparsity_dist.")
    sb = SparseBlob(model_blob)
    
    df = sb._rpt[:-1]
    df['tile_limit'] = df.tile_max - df.tile_min

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))  # 1 row, 2 columns

    plot_layer_hist(df, df.tile_limit.idxmin(), axes[0], type="narrowest")
    plot_layer_hist(df, df.tile_limit.idxmax(), axes[1], type="widest")

    plt.tight_layout()
    plt.savefig(f'{plot_name}.png')