script:
    '../scripts/select_bad_samples.py'

script:
    '../scripts/plink_filter.py'

script:
    '../scripts/pre_imputation_check.py'

shell:
    """
    plink --bfile {params.input} --a1-allele plink/{wildcards.batch}_merged_filter.bim.force_allele --make-bed --out plink/{wildcards.batch}_merged_mapped_alleled |& tee -a {log.plink}
    plink --bfile plink/{wildcards.batch}_merged_mapped_alleled --keep-allele-order --output-chr M --export vcf bgz --out vcf/{wildcards.batch}_merged_mapped_clean |& tee -a {log.vcf}
    mkdir vcf/temp_{wildcards.batch}
    bcftools sort -T vcf/temp_{wildcards.batch} vcf/{wildcards.batch}_merged_mapped_clean.vcf.gz -O z -o {output.temp_vcf} |& tee -a {log.vcf}
    bcftools index -f {output.temp_vcf} |& tee -a {log.vcf}
    # need to check output for the potential issues
    bcftools view {output.temp_vcf} --regions 1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22 -O b -o {output.bcf}
    bcftools norm --check-ref e -f {GRCH37_FASTA} vcf/{wildcards.batch}_merged_mapped_sorted.bcf.gz -O u -o /dev/null |& tee -a {log.vcf}
    bcftools index -f {output.bcf} | tee -a {log.vcf}
    """

shell:
    '''
        minimac4 \
        --refHaps {input.refHaps} \
        --haps phase/{wildcards.batch}_chr{wildcards.chrom}.phased.vcf.gz \
        --format GT,GP \
        --prefix imputed/{wildcards.batch}_chr{wildcards.chrom}.imputed \
        --minRatio 0.01 |& tee {log}
    '''

shell:
    '''
    FILTER="'strlen(REF)=1 & strlen(ALT)=1'"

    bcftools view -i 'strlen(REF)=1 & strlen(ALT)=1' imputed/{wildcards.batch}_chr{wildcards.chrom}.imputed.dose.vcf.gz -v snps -m 2 -M 2 -O z -o imputed/{wildcards.batch}_chr{wildcards.chrom}.imputed.dose.pass.vcf.gz |& tee {log}
    '''

shell:
    '''
    # for now just skip empty files
    true > {params.list} && \
    for i in {input}; do
        if [ -s $i ]
        then
            echo $i >> {params.list}
        else
            continue
        fi
    done

    bcftools concat -f {params.list} -O z -o {output} |& tee -a {log}
    bcftools index -f {output} |& tee -a {log}

    # check if there is a background data and merge it
    if [ -f "{params.merged_imputed}" && {params.mode} = "client" ]; then
        mv {output} {output}.client
        bcftools merge --force-samples {params.merged_imputed} {output}.client -O z -o {output} |& tee -a {log}
        bcftools index -f {output} |& tee -a {log}
    fi
    '''

shell:
    '''
    plink --vcf {input} --make-bed --out {params.out} |& tee {log}
    '''

shell:
    '''
    # please mind a merge step in merge_imputation_filter for germline
    plink --vcf {input} --make-bed --out {params.out} | tee {log}
    plink --bfile {params.background} --bmerge {params.out} --make-bed --out plink/{wildcards.batch}_merged_imputed |& tee {log}
    '''

shell:
    '''
        bcftools index -f -t {input.bcf} |& tee {log}
    '''

shell:
    '''
    # for now just skip empty files
    true > {params.list} && \
    for i in {input}; do
        if [ -s $i ]
        then
            echo $i >> {params.list}
        else
            continue
        fi
    done

    bcftools concat -f {params.list} -O z -o {output} |& tee -a {log}
    bcftools index -f {output} |& tee -a {log}

    # check if there is a background data and merge it
    if [ -f "background/{wildcards.batch}_merged_imputed.vcf.gz" && {params.mode} = "client" ]; then
        mv {output} {output}.client
        bcftools merge --force-samples background/{wildcards.batch}_merged_imputed.vcf.gz {output}.client -O b -o {output} |& tee -a {log}
        bcftools index -f {output} |& tee -a {log}
    fi
    '''

shell:
    '''
    bcftools query --list-samples input.vcf.gz >> vcf/samples.txt
    total_lines=$(wc -l < vcf/samples.txt)
    num_files={params.num_batches}
    ((lines_per_file = (total_lines + num_files - 1) / num_files))
    split -l $lines_per_file vcf/samples.txt vcf/batch --additional-suffix=.txt --numeric-suffixes=1
    for file in $(find vcf -name 'batch0[1-9].txt')
    do
        new=$(echo "$file" | sed "s/0//g")
        mv "$file" "$new"
    done
    '''

shell:
    '''
    bcftools view -S {input.samples} {input.vcf} -O z -o {output.vcf} --force-samples
    '''

shell:
    '''
        ln -sr {input.vcf} {output.vcf}
    '''

script: '../scripts/remove_imputation.py'

shell:
    '''
        ln -sr {input.vcf} {output.vcf}
    '''

shell:
    '''
        ln -sr {input.vcf} {output.vcf}
    '''

shell:
    '''
        bcftools annotate --set-id "%CHROM:%POS:%REF:%FIRST_ALT" {input.vcf} | \
        bcftools view -t ^8:10428647-13469693,21:16344186-19375168,10:44555093-53240188,22:16051881-25095451,2:85304243-99558013,1:118434520-153401108,15:20060673-25145260,17:77186666-78417478,15:27115823-30295750,17:59518083-64970531,2:132695025-141442636,16:19393068-24031556,2:192352906-198110229 | \
        bcftools view --min-af 0.05 --types snps -O b -o {output.bcf}
    '''

shell:
    '''
        ln -sr {input.bcf} {output.bcf}
    '''

shell:
    '''
        bcftools view {input.bcf} -O z -o {output.vcf}
    '''

shell:
    '''
    plink --vcf {input} --make-bed --out {params.out} |& tee {log}
    '''

shell:
    '''
    rm files_list.txt || true
    for file in {input}
    do
        if [[ $file == *.fam ]]
        then
            echo ${{file%.*}} >> files_list.txt
        fi
    done

    plink --merge-list files_list.txt --make-bed --out preprocessed/data
    rm files_list.txt
    '''

shell:
    '''
    bcftools index -f {input.batches_vcf}
    '''

shell:
    '''
    rm complete_vcf_list.txt || true
    for FILE in {input}
    do
        if [[ $FILE == *.gz ]]
        then
            echo $FILE >> complete_vcf_list.txt
        fi
    done
    bcftools merge --threads {threads} --file-list complete_vcf_list.txt --force-samples -O z -o {output.vcf}
    rm complete_vcf_list.txt
    '''

shell:
    '''
    plink --vcf {input} --make-bed --out {params.out} |& tee {log}
    '''

shell:
    '''
    (add-map-plink.pl -cm {input.bim} {params.genetic_map_GRCh37}> {output}) |& tee -a {log}
    '''

shell:
    '''
        bcftools query --list-samples {input.bcf} >> {output.fam}
    '''

shell:
    """
    ibis {input.bed} {input.bim} {input.fam} -t {threads} -mt {params.mT} -mL {params.mL} -ibd2 -mL2 3 -hbd -f ibis/merged_ibis |& tee -a {log}
    """

shell:
    """
    python {input.script} \
        --input-ibd-segments-file {input.ibd} \
        --mask-file {params.mask} \
        --output-ibd-segments-file {output.ibd}
    """

script:
    "../scripts/transform_ibis_segments.py"

shell:
    """
    zcat {input.ibd} > {output.ibd}
    ersa --avuncular-adj -ci -a {params.a} --dmax 14 -t {params.t} -l {params.l} \
            {params.r} -th {params.th} {output.ibd} -o {output.relatives} |& tee {log}
    """

shell:
    """
    FILES="{input}"
    for input_file in $FILES; do
        if [[ "$input_file" == "${{FILES[0]}}" ]]; then
            cat $input_file >> {output}
        else
            sed 1d $input_file >> {output}
        fi
    done
    """

script:
    "../scripts/split_map.py"

script: "../scripts/merge_king_ersa.py"

shell:
    """
        plink --bfile {params.bfile} --chr {params.chr} --make-bed --out ibis_data/{params.chr} --threads {threads}
    """

shell:
    '''
    (add-map-plink.pl -cm {input.bim} {params.genetic_map_GRCh37}> {output}) |& tee -a {log}
    '''

shell:
    """
        ibis {input.bed} {input.bim} {input.fam} -t {threads} -mt {params.mT} -mL {params.mL} -ibd2 -mL2 3 -hbd -gzip -f ibis/merged_ibis_{params.chr} |& tee -a {log}
    """

shell:
    """
        zcat {input} >> {output.ibd}
    """

shell:
    """
    python {input.script} \
        --input-ibd-segments-file {input.ibd} \
        --mask-file {params.mask} \
        --output-ibd-segments-file {output.ibd}
    """

script:
    "../scripts/transform_ibis_segments.py"

shell:
    """
    zcat {input.ibd} > {output.ibd}
    ersa --avuncular-adj -ci -a {params.a} --dmax 14 -t {params.t} -l {params.l} \
            {params.r} -th {params.th} {output.ibd} -o {output.relatives} |& tee {log}
    """

shell:
    """
    FILES="{input}"
    for input_file in $FILES; do
        if [[ "$input_file" == "${{FILES[0]}}" ]]; then
            cat $input_file >> {output}
        else
            sed 1d $input_file >> {output}
        fi
    done
    """

script: "../scripts/postprocess_ersa.py"

shell:
    """
    bcftools index {input} | tee {log}
    """

shell:
    """
    bcftools filter {input} -r {wildcards.chrom} -O z -o vcf/for_rapid_{wildcards.chrom}.vcf.gz | tee {log}
    """

script:
    "../scripts/estimate_rapid_params.py"

script:
    "../scripts/interpolate.py"

    script:
        "../scripts/interpolate.py"

'''
rule erase_dosages:
    input:
        vcf=rules.index_and_split.output[0]
    output:
        vcf='vcf/erased_chr{chrom}.vcf.gz'
    params:
        vcf='vcf/erased_chr{chrom}'
    conda:
        'bcftools'
    shell:
        "bcftools annotate -x 'fmt' {input.vcf} -O z -o {output.vcf}"
'''

shell:
    """
        rapidibd -i {input.vcf} -g {input.g} -d {params.min_cm_len} -o {params.output_folder}/chr{wildcards.chrom} \
        -w {params.window_size} -r {params.num_runs} -s {params.num_success}
    """

shell:
    """
        for file in {input}
        do
            zcat ${{file}} >> {output}
        done
    """

script:
    "../scripts/transform_rapid_segments.py"

shell:
    """
    touch {output}
    FILES="{input.ibd}"
    TEMPFILE=ersa/temp_relatives.tsv
    rm -f $TEMPFILE
    rm -f {output}

    for input_file in $FILES; do
        ersa --avuncular-adj -ci -a {params.a} --dmax 14 -t {params.t} -l {params.l} \
            {params.r} -th {params.th} $input_file -o $TEMPFILE |& tee {log}

        if [[ "$input_file" == "${{FILES[0]}}" ]]; then
            cat $TEMPFILE >> {output}
        else
            sed 1d $TEMPFILE >> {output}
        fi
    done
    """

script: "../scripts/postprocess_ersa.py"

shell:
    """
    KING_DEGREE=3

    king -b {params.input}.bed --cpus {threads} --ibdseg --degree $KING_DEGREE --prefix {params.out} |& tee {log}
    king -b {params.input}.bed --cpus {threads} --kinship --degree $KING_DEGREE --prefix {params.kin} |& tee -a {log}

    # we need at least an empty file for the downstream analysis
    if [ ! -f "{output.king}" ]; then
        touch {output.king}
    fi
    if [ ! -f "{output.kinship}" ]; then
        touch {output.kinship}
    fi
    if [ ! -f "{output.kinship0}" ]; then
        touch {output.kinship0}
    fi
    if [ ! -f "{output.segments}" ]; then
        touch {output.segments}
    fi
    """

shell: "bcftools index -f {input}"

shell:
    """
    bcftools filter -r {wildcards.chrom} {input.vcf} | bcftools norm --rm-dup none -O z -o vcf/imputed_chr{wildcards.chrom}.vcf.gz |& tee {log}
    """

script:
    "../scripts/vcf_to_ped.py"

shell:
    """
    plink --file ped/imputed_chr{wildcards.chrom} --keep-allele-order --cm-map {input.cmmap} {wildcards.chrom} --recode --out cm/chr{wildcards.chrom}.cm |& tee {log}
    """

shell:
    """
    touch {output}
    germline -input ped/imputed_chr{wildcards.chrom}.ped cm/chr{wildcards.chrom}.cm.map -min_m 2.5 -err_hom 2 -err_het 1 -output germline/chr{wildcards.chrom}.germline |& tee {log}
    # TODO: germline returns some length in BP instead of cM - clean up is needed
    set +e
    grep -v MB germline/chr{wildcards.chrom}.germline.match > germline/chr{wildcards.chrom}.germline.match.clean && mv germline/chr{wildcards.chrom}.germline.match.clean germline/chr{wildcards.chrom}.germline.match
    set -e
    """

shell:
     "cat germline/*.match > {output}"

script:
    '../scripts/merge_ibd.py'

shell:
    """
    touch {output}
    ersa --avuncular-adj -ci --alpha {params.alpha} --dmax 14 -t {params.ersa_t} -l {params.ersa_l} -th {params.ersa_th} {input.ibd} -o {output} |& tee {log}
    """

script: "../scripts/merge_king_ersa.py"

from utils.ibd import merge_germline_segments, segments_to_germline, interpolate_all, read_germline_segments, read_pedsim_segments


if __name__ == '__main__':

    germline_path = snakemake.input['germline']
    map_dir = snakemake.params['cm_dir']
    gap = float(snakemake.params['merge_gap'])
    pedsim_path = snakemake.params['true_ibd']

    use_true_ibd = bool(snakemake.params['use_true_ibd'])
    if not use_true_ibd:
        segments = read_germline_segments(germline_path)
        segments = interpolate_all(segments, map_dir)
        segments = merge_germline_segments(segments, gap=gap)
        segments_to_germline(segments, snakemake.output['ibd'])
    else:
        true_segments = read_pedsim_segments(pedsim_path)
        segments_to_germline(true_segments, snakemake.output['ibd'])

import pandas
import numpy
import os
import logging
from collections import Counter, namedtuple
from utils.ibd import read_king_segments as rks, interpolate_all, Segment


def is_non_zero_file(fpath):
    return os.path.isfile(fpath) and os.path.getsize(fpath) > 0


def _sort_ids(data):
    unsorted_mask = data.id2 < data.id1
    data.loc[unsorted_mask, 'id1'], data.loc[unsorted_mask, 'id2'] = data.loc[unsorted_mask, 'id2'], data.loc[
        unsorted_mask, 'id1']


def read_bucket_dir(bucket_dir):

    total = None
    for file in os.listdir(bucket_dir):
        if not file.endswith('tsv'):
            continue
        path = os.path.join(bucket_dir, file)
        bucket = read_germline(path)
        if total is None:
            total = bucket
        else:
            total = total.append(bucket)
    return total


def read_germline(ibd_path):
    germline_names = [
        'fid_iid1',
        'fid_iid2',
        'chrom',
        'start_end_bp',
        'start_end_snp',
        'snp_len',
        'genetic_len',
        'len_units',
        'mismatches',
        'is_homozygous1',
        'is_homozygous2'
    ]

    data = pandas.read_table(ibd_path, header=None, names=germline_names)
    if data.shape[0] == 0:
        return pandas.DataFrame(columns=['id1', 'id2', 'seg_count_germline', 'total_seg_len_germline']).set_index(['id1', 'id2'])

    data.loc[:, 'id1'] = data.fid_iid1.str.split().str[1]
    data.loc[:, 'id2'] = data.fid_iid2.str.split().str[1]
    _sort_ids(data)
    segments_info = data.loc[:, ['id1', 'id2', 'chrom', 'genetic_len']].groupby(by=['id1', 'id2']).agg({'genetic_len': 'sum', 'chrom': 'count'})
    segments_info.rename({'genetic_len': 'total_seg_len_germline', 'chrom': 'seg_count_germline'}, axis=1, inplace=True)
    return segments_info


def map_king_degree(king_degree):
    degree_map = {
        'Dup/MZ': 0,
        'PO': 1,
        'FS': 2,
        '2nd': 2,
        '3rd': 3,
        '4th': 4,
        'NA': numpy.nan
    }
    return [degree_map[kd] for kd in king_degree]


def map_king_relation(king_degree):
    degree_map = {
        'Dup/MZ': 0,
        'PO': 'PO',
        'FS': 'FS',
        '2nd': 2,
        '3rd': 3,
        '4th': 4,
        'NA': numpy.nan
    }
    return [degree_map[kd] for kd in king_degree]


def read_king(king_path):
    # FID1    ID1     FID2    ID2     MaxIBD1 MaxIBD2 IBD1Seg IBD2Seg PropIBD InfType
    try:
        data = pandas.read_table(king_path)
        data.loc[:, 'id1'] = data.FID1.astype(str) + '_' + data.ID1.astype(str)
        data.loc[:, 'id2'] = data.FID2.astype(str) + '_' + data.ID2.astype(str)
        _sort_ids(data)

        data.loc[:, 'king_degree'] = map_king_degree(data.InfType)
        data.loc[:, 'king_relation'] = map_king_relation(data.InfType)
        data.loc[:, 'king_degree'] = data.king_degree.astype(float).astype(pandas.Int32Dtype())
        data.rename({'PropIBD': 'shared_genome_proportion', 'IBD1Seg': 'king_ibd1_prop', 'IBD2Seg': 'king_ibd2_prop'}, axis='columns', inplace=True)
        indexed = data.loc[:, ['id1', 'id2', 'king_degree', 'king_relation', 'king_ibd1_prop', 'king_ibd2_prop', 'shared_genome_proportion']].\
            set_index(['id1', 'id2'])
        return indexed

    except pandas.errors.EmptyDataError:
        return pandas.DataFrame(data=None, index=None,
                                columns=['id1',
                                         'id2',
                                         'king_degree',
                                         'king_relation',
                                         'king_ibd1_prop',
                                         'king_ibd2_prop',
                                         'shared_genome_proportion']).set_index(['id1', 'id2'])


def read_king_segments_chunked(king_segments_path, map_dir):
    total = None
    try:
        for i, chunk in enumerate(pandas.read_table(
                                                    king_segments_path, 
                                                    compression='gzip', 
                                                    dtype={'FID1': str, 'ID1': str, 'FID2': str, 'ID2': str}, 
                                                    chunksize=1e+6)):

            segments = {}
            for i, row in chunk.iterrows():
                id1 = row['FID1'] + '_' + row['ID1']
                id2 = row['FID2'] + '_' + row['ID2']
                seg = Segment(id1, id2, row['Chr'],
                            bp_start=row['StartMB']*1e+6, bp_end=row['StopMB']*1e+6)
                key = tuple(sorted((seg.id1, seg.id2)))
                if key not in segments:
                    segments[key] = [seg]
                else:
                    segments[key].append(seg)

            segments = interpolate_all(segments, map_dir)
            data = pandas.DataFrame(columns=['id1', 'id2', 'total_seg_len_king', 'seg_count_king'])
            logging.info(f'loaded and interpolated segments from chunk {i} for {len(segments)} pairs')
            rows = []
            for key, segs in segments.items():
                row = {'id1': key[0],
                    'id2': key[1],
                    'total_seg_len_king': sum([s.cm_len for s in segs]),
                    'seg_count_king': len(segs)}
                rows.append(row)

            rows_frame = pandas.DataFrame.from_records(rows, columns=['id1', 'id2', 'total_seg_len_king', 'seg_count_king'])
            data = pandas.concat([data, rows_frame], ignore_index=True)

            _sort_ids(data)
            data = data.set_index(['id1', 'id2'])
            total = total.append(data) if total is not None else data

        return total

    except pandas.errors.EmptyDataError:
        return pandas.DataFrame(columns=['id1', 'id2', 'total_seg_len_king', 'seg_count_king']).set_index(['id1', 'id2'])


def read_king_segments(king_segments_path, map_dir):
    try:
        segments = rks(king_segments_path)
        segments = interpolate_all(segments, map_dir)
        data = pandas.DataFrame(columns=['id1', 'id2', 'total_seg_len_king', 'seg_count_king'])
        logging.info(f'loaded and interpolated segments for {len(segments)} pairs')
        for key, segs in segments.items():
            row = {'id1': key[0],
                   'id2': key[1],
                   'total_seg_len_king': sum([s.cm_len for s in segs]),
                   'seg_count_king': len(segs)}
            data = data.append(row, ignore_index=True)

        _sort_ids(data)
        return data.set_index(['id1', 'id2'])

    except pandas.errors.EmptyDataError:
        return pandas.DataFrame(columns=['id1', 'id2', 'total_seg_len_king', 'seg_count_king']).set_index(['id1', 'id2'])


def kinship_to_degree(kinship):
    degrees = []
    # intervals are from KING manual
    # >0.354, [0.177, 0.354], [0.0884, 0.177] and [0.0442, 0.0884]
    for k in kinship:
        if k > 0.354:
            degrees.append(0)
        elif 0.177 < k <= 0.354:
            degrees.append(1)
        elif 0.0884 < k <= 0.177:
            degrees.append(2)
        elif 0.0442 < k <= 0.0884:
            degrees.append(3)
        else:
            degrees.append(None)
    return degrees


def _read_kinship_data(kinship_path):
    try:
        data = pandas.read_table(kinship_path)
        print(kinship_path, data.shape)
        # If no relations were found, king creates file with only header
        if data.shape[0] != 0:
            if 'FID' in data.columns:
                data.loc[:, 'id1'] = data.FID.astype(str) + '_' + data.ID1.astype(str)
                data.loc[:, 'id2'] = data.FID.astype(str) + '_' + data.ID2.astype(str)
            else:
                data.loc[:, 'id1'] = data.FID1.astype(str) + '_' + data.ID1.astype(str)
                data.loc[:, 'id2'] = data.FID2.astype(str) + '_' + data.ID2.astype(str)

            _sort_ids(data)
            data.rename({'Kinship': 'kinship'}, axis=1, inplace=True)
            data.loc[:, 'kinship_degree'] = kinship_to_degree(data.kinship)
            data = data.loc[:, ['id1', 'id2', 'kinship', 'kinship_degree']].set_index(['id1', 'id2'])
        else:
            data = pandas.DataFrame(columns=['id1', 'id2', 'kinship', 'kinship_degree']).set_index(['id1', 'id2'])
        return data
    except pandas.errors.EmptyDataError:
        return pandas.DataFrame(columns=['id1', 'id2', 'kinship', 'kinship_degree']).set_index(['id1', 'id2'])


def read_kinship(kinship_path, kinship0_path):
    # parse within families
    # FID     ID1     ID2     N_SNP   Z0      Phi     HetHet  IBS0    Kinship Error
    within = _read_kinship_data(kinship_path)
    logging.info(f'loaded {within.shape[0]} pairs from within-families kinship estimation results')

    # FID1    ID1     FID2    ID2     N_SNP   HetHet  IBS0    Kinship
    across = _read_kinship_data(kinship0_path)
    logging.info(f'loaded {across.shape[0]} pairs from across-families kinship estimation results')
    return pandas.concat([within, across], axis=0)


def read_ersa(ersa_path):
    # Indv_1     Indv_2      Rel_est1      Rel_est2      d_est     lower_d  upper_d     N_seg     Tot_cM
    data = pandas.read_table(ersa_path, header=0,
                             names=['id1', 'id2', 'rel_est1', 'rel_est2',
                                    'ersa_degree', 'ersa_lower_bound', 'ersa_upper_bound', 'seg_count_ersa', 'total_seg_len_ersa'],
                             dtype={'ersa_degree': str, 'seg_count_ersa': str, 'total_seg_len_ersa': str})

    data = data.loc[(data.rel_est1 != 'NA') | (data.rel_est2 != 'NA'), :]
    data.loc[:, 'id1'] = data.id1.str.strip()
    data.loc[:, 'id2'] = data.id2.str.strip()
    _sort_ids(data)
    data.loc[:, 'ersa_degree'] = pandas.to_numeric(data.ersa_degree.str.strip(), errors='coerce').astype(
        pandas.Int32Dtype())
    data.loc[:, 'ersa_lower_bound'] = pandas.to_numeric(data.ersa_lower_bound.str.strip(), errors='coerce').astype(
        pandas.Int32Dtype())
    data.loc[:, 'ersa_upper_bound'] = pandas.to_numeric(data.ersa_upper_bound.str.strip(), errors='coerce').astype(
        pandas.Int32Dtype())
    print('dtypes are: ', data.dtypes)
    data.loc[:, 'seg_count_ersa'] = pandas.to_numeric(data.seg_count_ersa.str.strip(), errors='coerce').astype(
        pandas.Int32Dtype())
    data.loc[:, 'total_seg_len_ersa'] = pandas.to_numeric(data.total_seg_len_ersa.str.strip(), errors='coerce')

    data.loc[data.ersa_lower_bound != 1, 'ersa_lower_bound'] = data.ersa_lower_bound - 1
    data.loc[data.ersa_upper_bound != 1, 'ersa_upper_bound'] = data.ersa_upper_bound - 1

    logging.info(f'read {data.shape[0]} pairs from ersa output')
    logging.info(f'ersa after reading has {pandas.notna(data.ersa_degree).sum()}')
    data.loc[:, 'is_niece_aunt'] = [True if 'Niece' in d else False for d in data.rel_est1]
    logging.info(f'we have total of {data.is_niece_aunt.sum()} possible Niece/Aunt relationships')

    return data.loc[data.id1 != data.id2, ['id1', 'id2', 'ersa_degree', 'ersa_lower_bound', 'ersa_upper_bound', 'is_niece_aunt', 'seg_count_ersa', 'total_seg_len_ersa']].\
        set_index(['id1', 'id2'])


def read_ersa2(ersa_path):
    # individual_1    individual_2    est_number_of_shared_ancestors  est_degree_of_relatedness       0.95 CI_2p_lower:4
    # 2p_upper:5     1p_lower:6 1p_upper:7        0p_lower:8        0p_upper:9
    # maxlnl_relatedness      maxlnl_unrelatednes
    data = pandas.read_table(ersa_path, comment='#')

    data = data.loc[data.est_degree_of_relatedness != 'no_sig_rel', :]
    data.rename({'individual_1': 'id1', 'individual_2': 'id2', 'est_number_of_shared_ancestors': 'shared_ancestors'}, axis=1, inplace=True)
    data.loc[:, 'ersa_degree'] = pandas.to_numeric(data['est_degree_of_relatedness'], errors='coerce').\
        astype(pandas.Int32Dtype())

    lower, upper = [], []
    for i, ancestors in enumerate(data.shared_ancestors):
        if ancestors == 0:
            lower.append(data.iloc[i, 8])
            upper.append(data.iloc[i, 9])
        if ancestors == 1:
            lower.append(data.iloc[i, 6])
            upper.append(data.iloc[i, 7])
        if ancestors == 2:
            lower.append(data.iloc[i, 4])
            upper.append(data.iloc[i, 5])

    data.loc[:, 'ersa_lower_bound'] = lower
    data.loc[:, 'ersa_upper_bound'] = upper

    cols = ['id1', 'id2', 'ersa_degree', 'ersa_lower_bound', 'ersa_upper_bound', 'shared_ancestors']
    return data.loc[data.id1 != data.id2, cols].set_index(['id1', 'id2'])


if __name__ == '__main__':

    try:
        snakemake
    except NameError:
        test_dir = '/media_ssd/pipeline_data/TF-CEU-TRIBES-ibis-king-2'
        Snakemake = namedtuple('Snakemake', ['input', 'output', 'params', 'log'])
        snakemake = Snakemake(
            input={'king': f'{test_dir}/king/data.seg',
                   'king_segments': f'{test_dir}/king/data.segments.gz',
                   'kinship': f'{test_dir}/king/data.kin',
                   'kinship0': f'{test_dir}/king/data.kin0',
                   'ersa': f'{test_dir}/ersa/relatives.tsv'},

            output=[f'test_data/merge_king_ersa.tsv'],
            params={'cm_dir': f'{test_dir}/cm'},
            log=['test_data/merge_king_ersa.log']
        )

    logging.basicConfig(filename=snakemake.log[0], level=logging.DEBUG, format='%(levelname)s:%(asctime)s %(message)s')

    king_path = snakemake.input['king']
    king_segments_path = snakemake.input['king_segments']
    # within families
    kinship_path = snakemake.input['kinship']
    # across families
    kinship0_path = snakemake.input['kinship0']
    ersa_path = snakemake.input['ersa']
    map_dir = snakemake.params['cm_dir']
    output_path = snakemake.output[0]

    king = read_king(king_path)
    kinship = read_kinship(kinship_path, kinship0_path)
    king_segments = read_king_segments_chunked(king_segments_path, map_dir)
    ersa = read_ersa(ersa_path)

    relatives = king.merge(kinship, how='outer', left_index=True, right_index=True).\
                     merge(ersa, how='outer', left_index=True, right_index=True).\
                     merge(king_segments, how='outer', left_index=True, right_index=True)

    prefer_ersa_mask = pandas.isnull(relatives.king_degree) | (relatives.king_degree > 3)
    relatives.loc[:, 'final_degree'] = relatives.king_degree
    # if king is unsure or king degree > 3 then we use ersa distant relatives estimation
    relatives.loc[prefer_ersa_mask, 'final_degree'] = relatives.ersa_degree
    logging.info(f'king is null or more than 3: {prefer_ersa_mask.sum()}')
    logging.info(f'ersa is not null: {pandas.notna(relatives.ersa_degree).sum()}')
    print(f'relatives columns are {relatives.columns}')
    if 'total_seg_len_king' in relatives.columns:
        relatives.loc[:, 'total_seg_len'] = relatives.total_seg_len_king*relatives.king_ibd1_prop
        relatives.loc[:, 'total_seg_len_ibd2'] = relatives.total_seg_len_king*relatives.king_ibd2_prop
        relatives.loc[:, 'seg_count'] = relatives.seg_count_king

    relatives.loc[prefer_ersa_mask, 'total_seg_len'] = relatives.total_seg_len_ersa
    relatives.loc[prefer_ersa_mask, 'seg_count'] = relatives.seg_count_ersa
    print('is na: ', pandas.isna(relatives.loc[prefer_ersa_mask, 'total_seg_len_ersa']).sum())
    # approximate calculations, IBD share is really small in this case
    relatives.loc[prefer_ersa_mask, 'shared_genome_proportion'] = 0.5*relatives.loc[prefer_ersa_mask, 'total_seg_len'].values / 3440
    relatives.drop(['total_seg_len_king', 'seg_count_king', 'total_seg_len_ersa', 'seg_count_ersa', 'king_ibd1_prop', 'king_ibd2_prop'],
                   axis='columns', inplace=True)

    logging.info(f'final degree not null: {pandas.notna(relatives.final_degree).sum()}')

    relatives.loc[pandas.notna(relatives.final_degree), :].to_csv(output_path, sep='\t')