#!/usr/bin/env python

import argparse
from cosmosis.utils import read_chain_header,  extract_inis_from_chain_header

parser = argparse.ArgumentParser("Extract the input parameter files that were used to generate a cosmosis chain from the output.")
parser.add_argument("chain", help="Name of the chain file to read")
parser.add_argument("prefix", help="Prefix for output files {prefix}_params.ini, {prefix}_values.ini, {prefix}_priors.ini")



def save(lines, section, prefix):
    filename = "{}_{}.ini".format(prefix, section)
    open(filename,'w').writelines(lines)

def main(chain, prefix):
    lines = read_chain_header(chain)

    for section in ['params', 'values', 'priors']:
        section_lines = extract_inis_from_chain_header(lines, section)
        save(section_lines, section, prefix)

if __name__ == '__main__':
    args = parser.parse_args()
    main(args.chain, args.prefix)

