import sys, pysam, pickle

bam_path = "reads.bam"
offsets = []

# First pass: store file offsets for each read
with pysam.AlignmentFile(bam_path, "rb") as bam:
    for read in bam:
        offsets.append(bam.tell())

# Function to get a read by row index
def get_read_by_index(idx):
    with pysam.AlignmentFile(bam_path, "rb") as bam:
        bam.seek(offsets[idx])
        return next(bam)

# Example
read_10 = get_read_by_index(10)
print(read_10.query_name)


def main():
    in_file = sys.argv[1]
    out_file = sys.argv[2]

    with pysam.AlignmentFile(bam_path, "rb") as bam, open(out_file, "wb") as out:
        offsets = []
        for read in bam:
            offsets.append(bam.tell())
        pickle.dump(offsets, out)

if __name__=="__main__":
    main()