#!/usr/bin/env python

import argparse
import twixtools
import ismrmrd
import sys

defaults = {'out_file': 'merged_rawdata.dat'}

def parse_arguments(argv=None):
    parser = argparse.ArgumentParser(
        description='Merge Pulseq data into Siemens raw data file.',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument('in_file_1', type=str, help="Siemens raw data file acquired with original sequence.")
    parser.add_argument('in_file_2', type=str, help="Pulseq raw data file.")
    parser.add_argument('-o', '--out_file', type=str, help='Output Siemens raw data file name.')

    parser.set_defaults(**defaults)
    return parser.parse_args(argv)

def main(argv=None):
    args = parse_arguments(argv)

    if not args.in_file_1.endswith('.dat'):
        raise ValueError("Input file 1 must be a Siemens raw data file '.dat'.")
    if not (args.in_file_2.endswith('.dat') or args.in_file_2.endswith('.mrd')):
        raise ValueError("Input file 2 (Pulseq data file) must be either a Siemens raw data file '.dat' or ISMRMRD file '.mrd'.")
    use_mrd = args.in_file_2.endswith('.mrd')

    file1 = twixtools.read_twix(args.in_file_1, keep_syncdata=True, keep_acqend=True)
    n_acq1 = len(file1[-1]['mdb'])
    if use_mrd:
        file2 = ismrmrd.Dataset(args.in_file_2)
        n_acq2 = file2.number_of_acquisitions()
        if n_acq1 - n_acq2 == 1:
            n_acq2 += 1  # ACQEND is usually not converted to MRD
    else:
        file2 = twixtools.read_twix(args.in_file_2, parse_pmu=False, keep_syncdata=False, keep_acqend=True)
        n_acq2 = len(file2[-1]['mdb'])

    if n_acq1 != n_acq2:
        print(f"WARNING: Files have different number of measurement data blocks. File 1: {n_acq1}, File 2: {n_acq2}.")

    print(f"Read and copy {n_acq1} measurement data blocks.")
    for k, mdb in enumerate(file1[-1]['mdb']):
        if not mdb.is_flag_set('ACQEND') and not mdb.is_flag_set('SYNCDATA'):
            mdb = mdb.convert_to_local()
            if use_mrd:
                mdb.data = file2.read_acquisition(k).data.copy()
            else:
                mdb.data = file2[-1]['mdb'][k].data.copy()
            file1[-1]['mdb'][k] = mdb

    if use_mrd:
        file2.close()

    print(f"Writing merged data to {args.out_file}.")
    twixtools.write_twix(file1, args.out_file, version_is_ve=True)

if __name__ == "__main__":
    sys.exit(main())
