#!/usr/bin/env python

import optparse
import sys

import argweaver
import argweaver.vis

from rasmus import util


o = optparse.OptionParser(
    usage="%prog LAYOUT_FILE SMC_FILE",
    description="""
Write out layout for ARG stored in *.smc format
""")

conf, args = o.parse_args()


def layout_arg(smc_file, out):
    smc = argweaver.SMCReader(smc_file, parse_trees=True)
    names = smc.header["names"]
    chrom = smc.header["chrom"]

    for block, tree_layout in argweaver.vis.iter_layout_smc(smc, names=names):
        # Convert back to 1 index
        row = [chrom, block[0] + 1, block[1]]
        for name in names:
            row.extend([name, tree_layout[name]])
        util.print_row(*row, out=out)


# Parse arguments
if len(args) < 2:
    o.print_help()
    sys.exit(1)

layout_file = args[0]
smc_files = args[1:]

with argweaver.open_stream(layout_file, 'w', compress='bgzip') as out:
    for smc_file in smc_files:
        layout_arg(smc_file, out)
if layout_file.endswith('.gz'):
    argweaver.vis.index_arg_layout(layout_file)
