import h5py
import sys
import numpy as np
import os

def create_virtual_file( vsource, pro_indxs, datasource_shape, mydtype, filename   ):
    layout = h5py.VirtualLayout(shape=  tuple([ len(  pro_indxs   )     ]+ list(datasource_shape)[1:]), dtype=mydtype)
    for n in range(  len(  pro_indxs   )     ):
        layout[n] = vsource[ pro_indxs[n]   ]
    with h5py.File(filename, "w", libver="latest") as f:
        f.create_virtual_dataset("data", layout=layout, fillvalue=-5)


fn = sys.argv[1]
dn = sys.argv[2]
rep = sys.argv[3]

pyhst2_input = """
"""

entry = h5py.File(fn,"r")[dn]

pro_indxs = []
ff_indxs  = []
ff_poss   = []
bck_indxs = []

keys = entry["instrument/detector"]["image_key"]


old_ff_pos = -1
ffblock = None
ipro = 0
for i,k in enumerate(list(keys)):
    if k== 2:
        bck_indxs.append(i)
    if k== 0:
        pro_indxs.append(i)
        ipro += 1
    if k==1:
        if ipro== old_ff_pos:
            ffblock.append(i)
        else:
            ffblock = []
            ff_indxs.append( ffblock )
            ff_poss.append(ipro )
            old_ff_pos = ipro
        
datasource_shape =  entry["instrument/detector/data"] .shape
mydtype =  entry["instrument/detector/data"].dtype

vsource = h5py.VirtualSource(fn, dn+"/instrument/detector/data", shape= datasource_shape     )


filename = os.path.join(rep,"pyhst2_projs.h5")
create_virtual_file( vsource, pro_indxs, datasource_shape, mydtype, filename   )
pyhst2_input += """
FILE_PREFIX={}
PROJ_DS_NAME = "data"
""".format(  os.path.join(rep,"pyhst2_projs.h5")   )


if len(bck_indxs):
    filename = os.path.join(rep,"pyhst2_background.h5")
    create_virtual_file( vsource, bck_indxs, datasource_shape, mydtype, filename   )
    pyhst2_input += """
BACKGROUND_FILE={}
BACKGROUND_DS_NAME = "data"
    """.format(  os.path.join(rep,"pyhst2_projs.h5")   )


s_files="""["""
s_intervals="""["""
if len(ff_indxs):
    for ffblock, pos  in zip(ff_indxs, ff_pos  ) :
        s_intervals += """ {} ,""".format(pos)
        fname = "pyhst2_ff%05d.h5"%ff_pos
        fname = os.path.join(rep, fname)
        s_files+="""  "{}",""".format(fname=
                                      create_virtual_file( vsource, ffblock, datasource_shape, mydtype, fname )   )
                                      
        pyhst2_input += """

FF_PREFIX={ffprefix}]
FF_DS_NAME = "data"
        FF_INTERVALS={intervals}]
    """.format(  ffprefix=s_files[:-1]  , intervals=s_intervals[:-1] )

print (pyhst2_input)
