<filename>hic_hiker/main.py
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import argparse
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt
from sklearn.neighbors.kde import KernelDensity
import os
import time

# from memory_profiler import profile

# @profile
def main():
    psr = argparse.ArgumentParser()
    psr.add_argument('workspace', help='workspace directory')
    psr.add_argument('--fasta', help='assembler output fasta file containing contigs')
    psr.add_argument('--assembly', help='contig layout file: .assembly generated by 3D-DNA')
    psr.add_argument('--agp', help='contig layout file: .agp generate by SALSA etc')
    psr.add_argument('--mnd', help='mnd.txt file of Hi-C contacts')
    psr.add_argument('--sam', help='sam file of Hi-C contacts')
    psr.add_argument('--bam', help='bam file of Hi-C contacts')

    psr.add_argument('-K', default=75000, help='threshold', type=int)
    psr.add_argument('--refsam', help='sam file')
    args = psr.parse_args()

    # print('sam', args.sam)
    print('path =', os.path.abspath(args.workspace), args.K, type(args.K))

    from . import contigs, load_3ddna, prob, hmm, load, benchmark, layout, figures

    # workdir = os.path.dirname(os.path.abspath(__file__))
    workdir = os.path.abspath(args.workspace) + '/'

    # input checking
    is_3ddna = args.mnd and args.assembly
    is_general = args.agp and (args.sam or args.bam)
    if is_3ddna and not is_general:
        print('3ddna mode')
    elif is_general and not is_3ddna:
        print('general mode')
    else:
        print('[error] input data is inconsistent.')
        return 1

    is_bam = args.bam and not args.sam

    if args.fasta:
        print('[step 1] Loading files')
        print('loading', args.fasta)
        print('time', time.time())
        contigs = contigs.Contigs(
                fasta_filename=args.fasta
                )
    else:
        print('[step 1][error] Fasta file is required.')
        return 1

    # load files
    if is_3ddna:
        print('loading', args.assembly)
        print('time', time.time())
        asm = load_3ddna.Assembly(
                asm_filename=args.assembly,
                contigs=contigs
                )
        print('loading', args.mnd)
        print('time', time.time())
        _df = load.get_contacts_mnd(
                contigs=contigs,
                mnd_filename=args.mnd
                )
        print('time', time.time())
        original_layout = asm.get_layout()
        _contigs = contigs
        contigs, df = asm.update_all(_contigs, _df)
        print('time', time.time())

    if is_general:
        print('loading', args.agp)
        print('time', time.time())
        original_layout = layout.get_layout_from_agp(
                args.agp,
                contigs
                )
        # convert to assembly file for benchmark
        layout.generate_assembly_from_layout(
                workdir + 'original.assembly',
                contigs,
                original_layout
                )

        if is_bam:
            print('loading', args.bam)
            print('time', time.time())
            df = load.get_contacts_paired_sam(
                    contigs,
                    args.bam,
                    is_bam=True
                    )
        else:
            print('loading', args.sam)
            print('time', time.time())
            df = load.get_contacts_paired_sam(
                    contigs,
                    args.sam,
                    is_bam=False
                    )
    # loaded
    # contigs, df, original_layout

    print('[step 2] estimating the contact probability distribution')
    estimator, raw_estimator = prob.infer_from_longest_contig2(
            df,
            contigs,
            K=args.K,
            K0=10**3.5
            )
    print('time', time.time())

    plt.figure()
    figures.fig_distribution(contigs, df, args.K)
    plt.savefig(workdir + 'fig_distribution.pdf', bbox_inches='tight', pad_inches=0.05)

    plt.figure()
    figures.fig_distribution(contigs, df, args.K, mode='log')
    plt.savefig(workdir + 'fig_distribution_log.pdf', bbox_inches='tight', pad_inches=0.05)

    print('[step 3] calculating emission probabilities')
    print('time', time.time())
    probs = prob.get_prob(
            df,
            contigs,
            original_layout,
            estimator,
            max_k=20,
            show_progress=True
            )

    print('[step 4] running optimization algorithm')
    print('time', time.time())
    polished_layouts = {}

    for k in [2,3,4,5]:
        print('k =', k)
        polished_layouts[k] = hmm.optimize_layout(
                probs,
                contigs,
                original_layout,
                k=k
                )

    print('adaptive')
    polished_layouts[1] = hmm.optimize_layout(
            probs,
            contigs,
            original_layout,
            K=args.K
            )

    print('[step 5] Writing FASTA and assembly')
    print('time', time.time())
    if is_3ddna:
        asm.generate_assembly_with_new_layout(
                workdir + 'polished.assembly',
                contigs,
                polished_layouts[1]
                )
        for k in [2,3,4,5]:
            asm.generate_assembly_with_new_layout(
                    workdir + 'polishedk{}.assembly'.format(k),
                    contigs,
                    polished_layouts[k]
                    )
    elif is_general:
        layout.generate_assembly_from_layout(
                workdir + 'polished.assembly',
                contigs,
                polished_layouts[1]
                )
    layout.get_fasta(
            polished_layouts[1],
            contigs,
            workdir + 'polished.fasta',
            add_gap=True,
            )

    if args.refsam:
        print('[step 6] benchmarking')
        print('time', time.time())
        ref_layout = layout.get_reference_layout_from_sam(
                args.refsam,
                contigs
                )
        results = {}
        results[0] = benchmark.determine_correct_orientation_or_not(contigs, original_layout, ref_layout)[0]
        for k in [1,2,3,4,5]:
            results[k] = benchmark.determine_correct_orientation_or_not(contigs, polished_layouts[k], ref_layout)[0]

        plt.figure()
        figures.fig_errorchart(results)
        plt.savefig(workdir + 'fig_errorchart.pdf', bbox_inches='tight', pad_inches=0.05)

        plt.figure()
        figures.fig_length_error(contigs, original_layout, results[0], results[1])
        plt.savefig(workdir + 'fig_length_error.pdf', bbox_inches='tight', pad_inches=0.05)

        plt.figure()
        figures.fig_matrix(probs, contigs, polished_layouts[3], original_layout, results[3])
        plt.savefig(workdir + 'fig_matrix.pdf', bbox_inches='tight', pad_inches=0.05)

    print('finished all process')
    print('time', time.time())

if __name__ == '__main__':
    main()
