#!/usr/bin/env python3
import argparse
import pickle
import sys

import numpy

from frea.formats import *
from frea.simulation import *

parser = argparse.ArgumentParser(description='Second pass marginal summary statistic computation for simulation')
parser.add_argument('chromosome', help='Chromosome code (BED format)')
parser.add_argument('pheno', help='Merged phenotype file (pickle file)')
parser.add_argument('events', help='Recombination events file (pickle file output by frea-sample-events)')
parser.add_argument('--chunk', default=1000, help='Number of univariate model to fit in parallel')
parser.add_argument('--covars', help='Covariate file')
parser.add_argument('--output-sumstats', action="store_true", default=False, help='Output sumstats (ldsc) format')
args = parser.parse_args()

with open(args.pheno, 'rb') as f:
    y = pickle.load(f)
with open(args.events, 'rb') as f:
    state = pickle.load(f)
if args.covars is not None:
    with open(args.covars, 'rb') as f:
        covars, _ = pickle.load(f)
else:
    covars = None
stats = marginal_association(y=state['y'], covars=covars, seed=state['seed'],
                             n=state['y'].shape[0], events=state['events'],
                             chunk_size=int(args.chunk),
                             hotspot_file=state['hotspot_file'],
                             legend_file=state['legend_file'],
                             haps_file=state['haps_file'],
                             sample_file=state['sample_file'])
if args.output_sumstats:
    print('SNP', 'A1', 'A2', 'N', 'P', 'Z', sep='\t')
for name, pos, a0, a1, beta, se, logp in stats:
    if args.output_sumstats:
        print(name, a0, a1, state['y'].shape[0], 10 ** -logp, beta / se, sep='\t')
    else:
        print(args.chromosome, pos, pos + len(a0),
              get_pouyak_name(args.chromosome, name, pos, a0, a1), logp, sep='\t')
