from pynwb import NWBFile, NWBHDF5IO
from datetime import datetime
from nwbext_simulation_output import CompartmentSeries, Compartments


CELLS = [
    (0, [0, 1, 2, 3, 4], [.1, .2, .3, .4, .5]),
    (1, [0, 1, 2, 3], [.1, .2, .3, .4]),
    (2, [0, 1, 2, 3, 4, 5], [.1, .2, .3, .4, .5, .6]),
    (3, [0, 1, 2, 3, 4, 5], [.1, .2, .3, .4, .5, .6]),
]

BUFFER_SIZE = 10

if __name__ == '__main__':

    recorder = cell_vars.CellVarRecorderNWB('test.nwb', '~/src/bmtk/tmp', ['vm', 'im'])

    for gid, sec_list, seg_list in CELLS:
        recorder.add_cell(gid, sec_list, seg_list)
    
    recorder.tstop = 0.1
    recorder.initialize(100, buffer_size=BUFFER_SIZE)

    for i in range(100):
        for gid, sec_list, seg_list in CELLS:
            recorder.record_cell(gid, 'vm', [i]*len(seg_list), i)
            recorder.record_cell(gid, 'im', [i]*len(seg_list), i)

        if (i+1) % BUFFER_SIZE == 0:
            recorder.flush()
    
    recorder.close()