# visualize_heatmap.py
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import sys

def visualize_heatmap(txt_path:Path, out_path:Path = None, cmap: str = "inferno"):
    """Load a heat grid from TXT (first line: width height, then values) and visualize/save as PNG."""
    txt_path = Path(txt_path)
    with txt_path.open("r") as f:
        w, h = map(int, f.readline().split())
        data = np.fromstring(f.read(), sep=' ')
    grid = data.reshape((h, w))

    plt.figure(figsize=(6, 6), dpi=150)
    plt.imshow(grid, cmap=cmap)
    plt.axis('off')
    plt.tight_layout()

    plt.savefig(out_path, bbox_inches='tight', pad_inches=0)
    plt.show()
    plt.close()

if __name__ == "__main__":
    '''
        python visualize_heatmap.py <input_txt_file> <output_png_file>
    '''
    print(sys.argv[1],sys.argv[2])
    visualize_heatmap(sys.argv[1], sys.argv[2])