<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def heatmap(args):
 """ %prog heatmap input.npy genome.json Plot heatmap based on .npy data file. The .npy stores a square matrix with bins of genome, and cells inside the matrix represent number of links between bin i and bin j. The `genome.json` contains the offsets of each contig/chr so that we know where to draw boundary lines, or extract per contig/chromosome heatmap. """

    p = OptionParser(heatmap.__doc__)
    p.add_option("--resolution", default=500000, type="int",
                 help="Resolution when counting the links")
    p.add_option("--chr", help="Plot this contig/chr only")
    p.add_option("--nobreaks", default=False, action="store_true",
                 help="Do not plot breaks (esp. if contigs are small)")
    opts, args, iopts = p.set_image_options(args, figsize="10x10",
                                            style="white", cmap="coolwarm",
                                            format="png", dpi=120)

    if len(args) != 2:
        sys.exit(not p.print_help())

    npyfile, jsonfile = args
    contig = opts.chr
    # Load contig/chromosome starts and sizes
    header = json.loads(open(jsonfile).read())
    resolution = header.get("resolution", opts.resolution)
    logging.debug("Resolution set to {}".format(resolution))
    # Load the matrix
    A = np.load(npyfile)

    # Select specific submatrix
    if contig:
        contig_start = header["starts"][contig]
        contig_size = header["sizes"][contig]
        contig_end = contig_start + contig_size
        A = A[contig_start: contig_end, contig_start: contig_end]

    # Several concerns in practice:
    # The diagonal counts may be too strong, this can either be resolved by
    # masking them. Or perform a log transform on the entire heatmap.
    B = A.astype("float64")
    B += 1.0
    B = np.log(B)
    vmin, vmax = 1, 7
    B[B < vmin] = vmin
    B[B > vmax] = vmax
    print(B)
    logging.debug("Matrix log-transformation and thresholding ({}-{}) done"
                  .format(vmin, vmax))

    # Canvas
    fig = plt.figure(1, (iopts.w, iopts.h))
    root = fig.add_axes([0, 0, 1, 1])       # whole canvas
    ax = fig.add_axes([.05, .05, .9, .9])   # just the heatmap

    breaks = header["starts"].values()
    breaks += [header["total_bins"]]   # This is actually discarded
    breaks = sorted(breaks)[1:]
    if contig or opts.nobreaks:
        breaks = []
    plot_heatmap(ax, B, breaks, iopts, binsize=resolution)

    # Title
    pf = npyfile.rsplit(".", 1)[0]
    title = pf
    if contig:
        title += "-{}".format(contig)
    root.text(.5, .98, title, color="darkslategray", size=18,
              ha="center", va="center")

    normalize_axes(root)
    image_name = title + "." + iopts.format
    # macOS sometimes has way too verbose output
    logging.getLogger().setLevel(logging.CRITICAL)
    savefig(image_name, dpi=iopts.dpi, iopts=iopts)