import os
import h5py
import numpy as np
import time
import json

from bec_lib.core import RedisConnector, BECMessage


def load_data() -> tuple:
    """

    Returns
    """
    proj_nr = 180
    basedir = f"/Users/janwyzula/PSI/test_data_reduced/projection_{proj_nr:06d}/"

    metadata_name = f"/Users/janwyzula/PSI/test_data_reduced/projection_{proj_nr:06d}.json"
    with open(metadata_name) as file:
        metadata = json.load(file)

    filenames = [fname for fname in os.listdir(basedir) if fname.endswith(".h5")]
    filenames.sort()

    for ii, fname in enumerate(filenames):
        with h5py.File(os.path.join(basedir, fname), "r") as h5file:
            if ii == 0:
                q = h5file["q"][...].T.squeeze()
                norm_sum = h5file["norm_sum"][...]
                data = np.zeros((len(filenames), *h5file["I_all"][...].shape))
            data[ii, ...] = h5file["I_all"][...]

    return data, q, norm_sum, metadata


def _get_projection_keys(producer):
    keys = producer.keys("px_stream/projection_*")
    if not keys:
        return []
    return keys


def send_data(data, q, norm_sum, bec_producer, metadata, proj_nr) -> None:
    """"""
    start = time.time()

    keys = _get_projection_keys(bec_producer)
    pipe = bec_producer.pipeline()
    proj_numbers = set(key.decode().split("px_stream/projection_")[1].split("/")[0] for key in keys)
    if len(proj_numbers) > 5:
        for entry in sorted(proj_numbers)[0:-5]:
            for key in bec_producer.keys(f"px_stream/projection_{entry}/*"):
                bec_producer.delete(topic=key, pipe=pipe)
                print(f"Deleting {key}")

    # Add new data
    return_dict = {"metadata": metadata, "q": q, "norm_sum": norm_sum}
    msg = BECMessage.DeviceMessage(signals=return_dict).dumps()
    bec_producer.set_and_publish(f"px_stream/projection_{proj_nr}/metadata", msg=msg, pipe=pipe)
    return_dict = {"proj_nr": proj_nr}
    msg = BECMessage.DeviceMessage(signals=return_dict).dumps()
    bec_producer.set_and_publish(f"px_stream/proj_nr", msg=msg, pipe=pipe)
    pipe.execute()
    for line in range(data.shape[0]):
        return_dict = {"data": data[line, ...]}
        msg = BECMessage.DeviceMessage(signals=return_dict).dumps()
        print(f"Sending line {line}")
        bec_producer.rpush(topic=f"px_stream/projection_{proj_nr}/data", msgs=msg)
    print(f"Time to send {time.time()-start} seconds")
    print(f"Rate {data.shape[0]/(time.time()-start)} Hz")
    print(f"Data volume {data.nbytes/1e6} MB")


if __name__ == "__main__":
    data, q, norm_sum, metadata = load_data()
    bec_producer = RedisConnector(["localhost:6379"]).producer()
    proj_nr = 180
    while True:
        send_data(data, q, norm_sum, bec_producer, metadata, proj_nr=proj_nr)
        time.sleep(30)
        bec_producer.delete(topic=f"px_stream/projection_{proj_nr}/data:val")
        # proj_nr = proj_nr + 1