#!/usr/bin/env python

import argparse
import h5py
import numpy as np
from matplotlib import pyplot as plt


def get_bit_values(img):
    vmax = img.max()
    maxbit = int(vmax).bit_length()
    uniq = np.unique(img)
    res = []
    for bit in range(maxbit):
        bv = (1 << bit)
        if any(uniq & bv):
            res.append(bv)
    return res


def calc_square_plots(n_plots):
    nrows = int(round(np.sqrt(n_plots)))
    ncols = int(np.ceil(n_plots / nrows))
    return nrows, ncols


def make_reasons():
    # this should not be needed, but there is overlap in the reasons
    class Trict(dict):
        def __setitem__(self, key, value):
            if key not in self:
                super().__setitem__(key, [])
            super().__getitem__(key).append(value)

    reasons = Trict()

    reasons[5] = "additional pixel mask"

    for index, gain in enumerate((0, 1, 3)):
        reasons[index + 6] = f"std dev of gain {gain}"

    for hg0 in (False, True):
        for gain in range(4):
            reasons[gain + 4 * hg0] = f"wrong gain ({gain=})" if not hg0 else f"wrong gain ({gain=}, hg0)"

    res = {}
    for k, v in sorted(reasons.items()):
        v = " / ".join(v) # this should not be needed
        res[1 << k] = v

    return res


def print_dict(d):
    for i, (k, v) in enumerate(sorted(d.items())):
        print(i, k, v, sep="\t")





parser = argparse.ArgumentParser(description="Analyse the pixel mask of a pedestal file")
parser.add_argument("fname", help="Path to a pedestal file")
clargs = parser.parse_args()


with h5py.File(clargs.fname) as f:
    image = f["pixel_mask"][()]

values = get_bit_values(image)

n_plots = len(values) + 2
nrows, ncols = calc_square_plots(n_plots)
_fig, axs = plt.subplots(nrows, ncols, sharex=True, sharey=True)
axs = iter(axs.flatten())

def plot(img, title):
    ax = next(axs)
    ax.imshow(img, cmap="turbo")
    ax.set_title(title)
    print(f"{img.sum()} x {title}")


plot(image, "raw")

mask = image.astype(bool)
plot(mask, "mask")


reasons = make_reasons()
for val in values:
    selected = (image & val) == 0
    reason = reasons[val]
    plot(selected, f"{val} | {val:b} | {reason}")


print()
print_dict(reasons)


plt.tight_layout()
try:
    plt.show()
except KeyboardInterrupt:
    pass