#!/usr/bin/env python

import argparse

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 read_comment_section(filename):
    lines = []
    for line in open(filename):
        if not line.startswith('#'):
            break
        lines.append(line)
    return lines



def extract_section(lines, section):
    start = "## START_OF_{}_INI".format(section).upper()
    end = "## END_OF_{}_INI".format(section).upper()
    in_section = False
    output_lines = []
    for line in lines:
        if line.startswith(start):
            in_section = True
            continue
        elif line.startswith(end):
            break
        elif in_section:
            output_lines.append(line[3:])
    return output_lines

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

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

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

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

