#!/usr/bin/env python3

'''
Copyright (c) 2020 Children's Hospital of Philadelphia
Author: Li Fang (fangli2718@gmail.com)

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

'''

import os
import sys
import numpy as np
import math
import gzip
import argparse

import tk

from sklearn.mixture import GaussianMixture

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from matplotlib import cm
matplotlib.rcParams['font.sans-serif'] = "Arial"
matplotlib.rcParams['font.family'] = "sans-serif"

debug = 1

class Readinfo:
    def __init__(self, readname):

        self.readname = readname
        self.label = -1
        self.repeat_size1 = -1
        self.repeat_size2 = -1
        self.max_prob = 0
        self.proba_list = list()
        self.qc_passed = 1

tab  = '\t'
endl = '\n'

def parse_user_arguments():

    parser = argparse.ArgumentParser(description='Joint quantification of two adjacent tandem repeats from long-read amplicon sequencing data ')
    ### required arguments ###
    parser.add_argument('--in_fq',        required = True,  metavar = 'PATH',   type = str, help = 'input fastq file')
    parser.add_argument('--platform',     required = True,  metavar = 'STRING', type = str, help = 'three valid values: `ont`, `pacbio`, `consensus`')
    parser.add_argument('--ref_fasta',    required = True,  metavar = 'PATH',   type = str, help = 'reference genome sequence in FASTA format')
    parser.add_argument('--repeat1',      required = True,  metavar = 'chr:start:end:repeat_unit:max_size', type = str, help = 'first tandem repeat, coordinates are 1-based (e.g. chr4:3074877:3074933:CAG:200)')
    parser.add_argument('--repeat2',      required = True,  metavar = 'chr:start:end:repeat_unit:max_size', type = str, help = 'second tandem repeat, coordinates are 1-based (e.g. chr4:3074946:3074966:CCG:20)')
    parser.add_argument('--out_dir',      required = True,  metavar = 'PATH',   type = str, help = 'path to the output directory')
    parser.add_argument('--version',      action='version', version='%(prog)s 0.3.0')

    ### optional arguments ### ploidy
    parser.add_argument('--num_threads',  required = False, metavar = 'INT',    type = int, default = 1,  help = 'number of threads used by minimap2 (default: 1)')
    parser.add_argument('--minimap2',     required = False, metavar = 'PATH',   type = str, default = '', help = 'path to minimap2 (default: using environment default)')
    parser.add_argument('--ploidy',       required = False, metavar = 'INT',    type = int, default = 2,  help = 'ploidy of the sample (default: 2)')

    input_args = parser.parse_args()

    input_args.platform = input_args.platform.strip('`')
    input_args.platform = input_args.platform.strip('\'')
    input_args.platform = input_args.platform.strip('\"')

    valid_platform_list = ['ont', 'pacbio', 'consensus']
    input_args.platform = input_args.platform.lower()

    if input_args.platform not in valid_platform_list:
        tk.eprint('ERROR: --platform should be one of the three valid values: ont, pacbio, consensus !\n')
        sys.exit(1)

    if input_args.ploidy < 1:
        tk.eprint('ERROR: --ploidy must be >= 1 !\n')
        sys.exit(1)

    tk.check_input_file_exists(input_args.in_fq)
    tk.check_input_file_exists(input_args.ref_fasta)

    input_args.in_fq       = os.path.abspath(input_args.in_fq)
    input_args.ref_fasta   = os.path.abspath(input_args.ref_fasta)
    input_args.out_dir     = os.path.abspath(input_args.out_dir)

    if input_args.minimap2 == '':
        input_args.minimap2 = tk.find_executable_path('minimap2')
        if not input_args.minimap2:
            tk.eprint('ERROR! minimap2 was not found! Please supply the path to --minimap2')
            sys.exit(1)
        else:
            if os.path.exists(input_args.minimap2):
                tk.eprint('NOTICE: found path to minimap2: %s' % input_args.minimap2)
            else:
                tk.eprint('ERROR! minimap2 was not found! Please supply the path to --minimap2')

    tk.check_input_file_exists(input_args.minimap2)

    if len(input_args.repeat1.split(':')) != 5:
        tk.eprint('ERROR! --repeat1 should be in this format: chr:start:end:repeat_unit:max_size (e.g. chr4:3074876:3074933:CAG:200)')
        sys.exit(1)

    if len(input_args.repeat2.split(':')) != 5:
        tk.eprint('ERROR! --repeat2 should be in this format: chr:start:end:repeat_unit:max_size (e.g. chr4:3074876-3074933:CAG:200)')
        sys.exit(1)

    return input_args

def main():

    input_args = parse_user_arguments()
    ampRepeat_joint (input_args)

    return

def ampRepeat_joint (input_args):

    n_input_reads = tk.count_fastq(input_args.in_fq)
    if n_input_reads < input_args.ploidy:
        tk.eprint(f'ERROR: No enough reads for analysis. Ploidy was set to {input_args.ploidy} but there were only {n_input_reads} reads in the fastq file: {input_args.in_fq}\n')
        sys.exit(1)

    tk.create_dir(input_args.out_dir)

    repeat1 = Repeat().init_from_string(input_args.repeat1)
    repeat2 = Repeat().init_from_string(input_args.repeat2)

    if repeat1.chrom != repeat2.chrom:
        tk.eprint('ERROR: joint quantification only works with two nearby repeat regions. The two repeats in the config_file are in different chromosomes and there is no need to do joint quantification.\n')
        sys.exit(1)

    if repeat1.start > repeat2.start:
        repeat1, repeat2 = tk.switch_two_objects(repeat1, repeat2)

    repeat1.max_size += 10
    repeat2.max_size += 10
    max_anchor_len = 1000

    if repeat1.end + 100 < repeat2.start:
        tk.eprint('ERROR: joint quantification only works with two nearby repeat regions (distance < 100 bp). The two repeats are far away from each other and there is no need to do joint quantification.\n')
        sys.exit(1)

    in_fastq_prefix = os.path.splitext(os.path.split(input_args.in_fq)[1])[0]
    temp_out_dir = os.path.join(input_args.out_dir, '%s.AmpRepeat.temp' % in_fastq_prefix)
    tk.create_dir(temp_out_dir)

    repeat_chrom_seq = tk.read_one_chr_from_fasta_file(input_args.ref_fasta, repeat1.chrom)
    if len(repeat_chrom_seq) == 0:
        tk.eprint('ERROR: ref_fasta file: %s has no valid sequence!\n' % input_args.ref_fasta)
        sys.exit(1)

    initial_estimation = initial_estimate_repeat_size(input_args.minimap2, repeat_chrom_seq, input_args.in_fq, input_args.platform, input_args.num_threads, repeat1, repeat2, max_anchor_len, temp_out_dir)

    log_fn = 'repeatRepeat.log'
    log_f  = open(log_fn, 'w')
    log_f.write(initial_estimation.output_repeat2_boundaries())

    final_estimation = fine_tune_read_count(initial_estimation, input_args.in_fq, repeat_chrom_seq, repeat1, repeat2, input_args.minimap2, input_args.platform, input_args.num_threads, temp_out_dir)

    jointly_split_alleles_using_gmm (input_args.ploidy, repeat1, repeat2, final_estimation, input_args.in_fq, input_args.out_dir)

    tk.eprint('NOTICE: program finished. Output files are here: %s\n' % input_args.out_dir)

    log_f.close()
    return

def output_repeat_size_file(read_repeat_joint_count_dict, readinfo_dict, in_fastq_file, repeat_region_list, repeat_size_file, ploidy):

    repeat_size_fp = open(repeat_size_file, 'w')
    repeat_size_fp.write('##input_fastq=%s\n' % in_fastq_file)
    repeat_size_fp.write('#readname\t%s\t%s\tallele_id\tprobability\tQC\n' % (repeat_region_list[0].repeat_id, repeat_region_list[1].repeat_id))
    for readname in read_repeat_joint_count_dict:
        repeat_size1, repeat_size2 = read_repeat_joint_count_dict[readname]
        if readname in readinfo_dict:
            readinfo = readinfo_dict[readname]
            allele_id = readinfo.label + 1
            if readinfo.qc_passed:
                qc = 'PASS'
            else:
                qc = 'FAIL'
            max_prob = readinfo.max_prob
            repeat_size_fp.write('%s\t%d\t%d\t%d\t%.8f\t%s' % (readname, repeat_size1, repeat_size2, allele_id, max_prob, qc))
            for prob in readinfo.proba_list:
                repeat_size_fp.write('\t%.8f' % prob )
            repeat_size_fp.write('\n')

        else:
            allele_id = 'N.A.'
            max_prob = 'N.A.'
            qc = 'FAIL'
            repeat_size_fp.write('%s\t%d\t%d\t%s\t%s\t%s' % (readname, repeat_size1, repeat_size2, allele_id, max_prob, qc))
            for i in range(0, ploidy):
                repeat_size_fp.write('\tN.A.')
            repeat_size_fp.write('\n')
    repeat_size_fp.close()

    return

def fine_tune_read_count(initial_estimation, in_fastq_file, repeat_chrom_seq, repeat1, repeat2, minimap2, platform, num_threads, out_dir):

    assert repeat1.chrom == repeat2.chrom
    assert repeat1.start < repeat2.start

    repeat1.round1_max_size = 0
    repeat2.round1_max_size = 0
    repeat1.round1_min_size = repeat1.max_size
    repeat2.round1_min_size = repeat2.max_size
    for readname, value in initial_estimation.repeat1_count_range_dict.items():
        min_repeat_size, max_repeat_size = value
        if max_repeat_size > repeat1.round1_max_size:
            repeat1.round1_max_size = max_repeat_size
        if min_repeat_size < repeat1.round1_min_size:
            repeat1.round1_min_size = min_repeat_size

    for readname, value in initial_estimation.repeat2_count_range_dict.items():
        min_repeat_size, max_repeat_size = value
        if max_repeat_size > repeat2.round1_max_size:
            repeat2.round1_max_size = max_repeat_size
        if min_repeat_size < repeat2.round1_min_size:
            repeat2.round1_min_size = min_repeat_size

    if repeat1.round1_max_size > repeat1.max_size: repeat1.round1_max_size = repeat1.max_size
    if repeat2.round1_max_size > repeat2.max_size: repeat2.round1_max_size = repeat2.max_size

    tk.eprint('NOTICE: In round 1 estimation, repeat 1 (%s) is in the range of (%d, %d)' % (repeat1.repeat_unit, repeat1.round1_min_size, repeat1.round1_max_size))
    tk.eprint('NOTICE: In round 1 estimation, repeat 2 (%s) is in the range of (%d, %d)' % (repeat2.repeat_unit, repeat2.round1_min_size, repeat2.round1_max_size))

    fastq_dict = fastq_file_to_dict(in_fastq_file)

    round2_estimation = round2_estimation_of_repeat_size(initial_estimation, fastq_dict, repeat_chrom_seq, repeat1, repeat2, minimap2, platform, num_threads, out_dir)

    if round2_estimation.step_size1 > 1 and round2_estimation.step_size2 > 1:
        final_estimation = round3_estimation_of_repeat_size(initial_estimation, round2_estimation, fastq_dict, repeat_chrom_seq, repeat1, repeat2, minimap2, platform, num_threads, out_dir)
    else:
        final_estimation = round2_estimation

    return final_estimation

def round3_estimation_of_repeat_size(initial_estimation, round2_estimation, fastq_dict, repeat_chrom_seq, repeat1, repeat2, minimap2, platform, num_threads, out_dir):

    if len(round2_estimation.repeat1_count_dict) == 0 or len(round2_estimation.repeat2_count_dict) == 0:
        return RepeatSize()

    assert repeat1.chrom == repeat2.chrom
    assert repeat1.start < repeat2.start
    max_flanking_len = 1000

    buffer_size1 = round2_estimation.step_size1
    buffer_size2 = round2_estimation.step_size2
    tk.eprint(f'NOTICE: Fine-tuning repeat size. buffer_size1 = {buffer_size1}, buffer_size2 = {buffer_size2}')

    round2_size1_list = list()
    round2_size2_list = list()

    for readname in round2_estimation.repeat1_count_dict:
        if readname not in round2_estimation.repeat2_count_dict: continue
        size1 = round2_estimation.repeat1_count_dict[readname]
        size2 = round2_estimation.repeat2_count_dict[readname]
        round2_size1_list.append(size1)
        round2_size2_list.append(size2)

    min_size1 = int(min(round2_size1_list) - buffer_size1)
    max_size1 = int(max(round2_size1_list) + buffer_size1 + 2)
    if min_size1 < 0: min_size1 = 0
    min_size2 = int(min(round2_size2_list) - buffer_size2)
    max_size2 = int(max(round2_size2_list) + buffer_size2 + 2)
    if min_size2 < 0: min_size2 = 0

    ## round 3 alignment ##
    left_anchor_seq, mid_anchor_seq, right_anchor_seq = extract_anchor_seq_for_two_repeats (repeat_chrom_seq, repeat1, repeat2, max_flanking_len)
    left_anchor_len  = len(left_anchor_seq)
    mid_anchor_len   = len(mid_anchor_seq)

    preset = tk.get_preset_for_minimap2(platform)
    round3_paf_file =  os.path.join(out_dir, 'round3.paf')
    round3_paf_fp = open(round3_paf_file, 'w')
    round3_paf_fp.close()

    for repeat_count1 in range(min_size1, max_size1):
        for repeat_count2 in range(min_size2, max_size2):
            tmp_fastq_file = os.path.join(out_dir, '%d.%d.round3.fastq' % (repeat_count1, repeat_count2) )
            tmp_fastq_fp = open(tmp_fastq_file, 'w')
            n_tmp_reads = 0
            for readname in fastq_dict:
                if readname not in round2_estimation.repeat1_count_dict: continue
                if readname not in round2_estimation.repeat2_count_dict: continue
                size1 = round2_estimation.repeat1_count_dict[readname]
                size2 = round2_estimation.repeat2_count_dict[readname]
                if repeat_count1 < size1 - buffer_size1 or repeat_count1 >= size1 + buffer_size1: continue
                if repeat_count2 < size2 - buffer_size2 or repeat_count2 >= size2 + buffer_size2: continue
                r1min1, r1max1 = initial_estimation.repeat1_count_range_dict[readname]
                r1min2, r1max2 = initial_estimation.repeat2_count_range_dict[readname]
                if repeat_count1 < r1min1 or repeat_count1 >= r1max1: continue
                if repeat_count2 < r1min2 or repeat_count2 >= r1max2: continue

                tmp_fastq_fp.write(fastq_dict[readname])
                n_tmp_reads += 1
            tmp_fastq_fp.close()
            if n_tmp_reads == 0:
                tk.rm(tmp_fastq_file)
                continue
            tmp_ref_file = os.path.join(out_dir, '%d.%d.ref.fasta' % (repeat_count1, repeat_count2))
            build_fasta_template_for_two_repeats(left_anchor_seq, mid_anchor_seq, right_anchor_seq, repeat1, repeat2, repeat_count1, repeat_count2, tmp_ref_file)
            cmd = f'{minimap2} -A 2 -B 6 -c --eqx -t {num_threads} -x {preset} {tmp_ref_file} {tmp_fastq_file} >> {round3_paf_file} 2> /dev/null'
            tk.run_system_cmd(cmd)
            tk.rm(tmp_ref_file)
            tk.rm(tmp_fastq_file)

    round3_estimation = estimate_two_repeats_from_paf(round3_paf_file, left_anchor_len, mid_anchor_len, repeat1, repeat2)
    round3_estimation.step_size1 = 1
    round3_estimation.step_size2 = 1

    return round3_estimation

def choose_best_step_size(repeat, count_range_dict):

    max_len = 50
    max_step_size = int(max_len/repeat.repeat_unit_size)
    if max_step_size < 1: max_step_size = 1

    error_list = list()
    for readname in count_range_dict:
        a, b = count_range_dict[readname]
        error = b - a
        error_list.append(error)

    l = np.mean(error_list)

    count_list = list()
    for size in range(1, max_step_size+1):
        count = int(l/size)+1
        count += size * 2 + 2
        tup = (size, count)
        count_list.append(tup)

    count_list.sort(key = lambda x:x[1])

    return count_list[0][0]

def round2_estimation_of_repeat_size(initial_estimation, fastq_dict, repeat_chrom_seq, repeat1, repeat2, minimap2, platform, num_threads, out_dir):

    assert repeat1.chrom == repeat2.chrom
    assert repeat1.start < repeat2.start
    max_flanking_len = 1000

    step_size1 = choose_best_step_size(repeat1, initial_estimation.repeat1_count_range_dict)
    step_size2 = choose_best_step_size(repeat2, initial_estimation.repeat2_count_range_dict)

    tk.eprint(f'NOTICE: Round 2 estimation. step_size1 = {step_size1}; step_size2 = {step_size2}')

    ## round 2 alignment ##
    left_anchor_seq, mid_anchor_seq, right_anchor_seq = extract_anchor_seq_for_two_repeats (repeat_chrom_seq, repeat1, repeat2, max_flanking_len)
    left_anchor_len  = len(left_anchor_seq)
    mid_anchor_len   = len(mid_anchor_seq)

    preset = tk.get_preset_for_minimap2(platform)
    round2_paf_file =  os.path.join(out_dir, 'round2.paf')
    round2_paf_fp = open(round2_paf_file, 'w')
    round2_paf_fp.close()

    for repeat_count1 in range(repeat1.round1_min_size, repeat1.round1_max_size + 1, step_size1):
        for repeat_count2 in range(repeat2.round1_min_size, repeat2.round1_max_size + 1, step_size2):
            tmp_fastq_file = os.path.join(out_dir, '%d.%d.round2.fastq' % (repeat_count1, repeat_count2) )
            tmp_fastq_fp = open(tmp_fastq_file, 'w')
            n_tmp_reads = 0
            for readname in fastq_dict:
                if readname not in initial_estimation.repeat1_count_range_dict: continue
                if readname not in initial_estimation.repeat2_count_range_dict: continue
                min1, max1 = initial_estimation.repeat1_count_range_dict[readname]
                min2, max2 = initial_estimation.repeat2_count_range_dict[readname]
                if repeat_count1 >= min1 and repeat_count1 < max1 and repeat_count2 >= min2 and repeat_count2 < max2:
                    tmp_fastq_fp.write(fastq_dict[readname])
                    n_tmp_reads += 1
            tmp_fastq_fp.close()
            if n_tmp_reads == 0:
                tk.rm(tmp_fastq_file)
                continue
            tmp_ref_file = os.path.join(out_dir, '%d.%d.ref.fasta' % (repeat_count1, repeat_count2))
            build_fasta_template_for_two_repeats(left_anchor_seq, mid_anchor_seq, right_anchor_seq, repeat1, repeat2, repeat_count1, repeat_count2, tmp_ref_file)
            cmd = f'{minimap2} -A 2 -B 6 -c --eqx -t {num_threads} -x {preset} {tmp_ref_file} {tmp_fastq_file} >> {round2_paf_file} 2> /dev/null'
            tk.run_system_cmd(cmd)
            tk.rm(tmp_ref_file)
            tk.rm(tmp_fastq_file)

    round2_estimation = estimate_two_repeats_from_paf(round2_paf_file, left_anchor_len, mid_anchor_len, repeat1, repeat2)
    round2_estimation.step_size1 = step_size1
    round2_estimation.step_size2 = step_size2

    return round2_estimation

def estimate_two_repeats_from_paf(in_paf_file, left_anchor_len, mid_anchor_len, repeat1, repeat2):

    repeat_size_estimation = RepeatSize()

    in_paf_fp = open(in_paf_file, 'r')
    read_align_score_dict = dict()
    while 1:
        line = in_paf_fp.readline()
        if not line: break
        line = line.strip()
        if not line: continue

        col_list = line.strip().split('\t')
        paf = tk.PAF(col_list)
        repeat_size1, repeat_size2 = paf.tname.split('-')
        repeat_size1 = int(repeat_size1)
        repeat_size2 = int(repeat_size2)

        a = left_anchor_len - 10
        if a < 0: a = 0
        b = left_anchor_len + repeat1.repeat_unit_size * repeat_size1 + mid_anchor_len + repeat2.repeat_unit_size * repeat_size2 + 10
        if b > paf.tlen: b = paf.tlen
        repeat_region_align_score = tk.target_region_alignment_stats_from_cigar(paf.cigar, paf.tstart, paf.tend, a, b).score

        tup = (repeat_size1, repeat_size2, repeat_region_align_score, paf.align_score, paf.qlen)
        if paf.qname not in read_align_score_dict:
            read_align_score_dict[paf.qname] = list()
        read_align_score_dict[paf.qname].append(tup)

    in_paf_fp.close()

    for readname in read_align_score_dict:
        tup_list = read_align_score_dict[readname]
        tup_list.sort(key = lambda x:x[2], reverse = True)
        max_score = tup_list[0][2]
        max_score_size1_list = list()
        max_score_size2_list = list()
        for tup in tup_list:
            if tup[2] == max_score:
                size1 = tup[0]
                size2 = tup[1]
                max_score_size1_list.append(size1)
                max_score_size2_list.append(size2)
            else:
                break

        round2_repeat_size1 = np.mean(max_score_size1_list)
        round2_repeat_size2 = np.mean(max_score_size2_list)
        repeat_size_estimation.repeat1_count_dict[readname] = round2_repeat_size1
        repeat_size_estimation.repeat2_count_dict[readname] = round2_repeat_size2

    return repeat_size_estimation

def extract_anchor_seq_for_two_repeats (repeat_chrom_seq, repeat1, repeat2, max_flanking_len):

    left_end_pos = repeat1.start
    left_start_pos = left_end_pos - max_flanking_len
    if left_start_pos < 0: left_start_pos = 0

    mid_start_pos = repeat1.end
    mid_end_pos = repeat2.start

    right_start_pos = repeat2.end
    right_end_pos = right_start_pos + max_flanking_len
    if right_end_pos > len(repeat_chrom_seq): right_end_pos = len(repeat_chrom_seq)

    left_anchor_seq  = repeat_chrom_seq[left_start_pos:left_end_pos]
    mid_anchor_seq   = repeat_chrom_seq[mid_start_pos:mid_end_pos]
    right_anchor_seq = repeat_chrom_seq[right_start_pos:right_end_pos]

    return left_anchor_seq, mid_anchor_seq, right_anchor_seq

def build_fasta_template_for_two_repeats(left_anchor_seq, mid_anchor_seq, right_anchor_seq, repeat1, repeat2, repeat_count1, repeat_count2, tmp_ref_file):

    tmp_ref_fp = open(tmp_ref_file, 'w')
    tmp_ref_fp.write('>%d-%d\n' % (repeat_count1, repeat_count2))
    seq = left_anchor_seq + repeat1.repeat_unit * repeat_count1 + mid_anchor_seq + repeat2.repeat_unit * repeat_count2 + right_anchor_seq + '\n'
    tmp_ref_fp.write(seq)
    tmp_ref_fp.close()

    return

def initial_estimate_repeat_size(minimap2, repeat_chrom_seq, in_fastq_file, platform, num_threads, repeat1, repeat2, max_anchor_len, out_dir):

    tk.eprint('NOTICE: Round 1 estimation')
    assert repeat1.chrom == repeat2.chrom
    assert repeat1.start < repeat2.start

    left_anchor_seq, right_anchor_seq = tk.extract_anchor_sequence(repeat_chrom_seq, repeat1.start, repeat2.end, max_anchor_len)
    left_anchor_len  = len(left_anchor_seq)
    right_anchor_len = len(right_anchor_seq)

    ## make template fasta file ##
    left_template_fasta_file = os.path.join(out_dir, 'round1_templates.left_anchor.fasta')
    left_template_name  = 'left_anchor_%d_%d_%s' % (len(left_anchor_seq), repeat1.max_size, repeat1.repeat_unit)
    left_template_seq   = left_anchor_seq + repeat1.repeat_unit * repeat1.max_size

    left_template_fasta_fp = open(left_template_fasta_file, 'w')
    left_template_fasta_fp.write('>%s\n' % left_template_name)
    left_template_fasta_fp.write('%s\n' % left_template_seq)
    left_template_fasta_fp.close()

    right_template_fasta_file = os.path.join(out_dir, 'round1_templates.right_anchor.fasta')
    right_template_name = 'right_anchor_%d_%d_%s_revc' % (len(right_anchor_seq), repeat2.max_size, repeat2.repeat_unit)
    right_template_seq  = repeat2.repeat_unit * repeat2.max_size + right_anchor_seq
    right_template_seq  = tk.rev_comp(right_template_seq)

    right_template_fasta_fp = open(right_template_fasta_file, 'w')
    right_template_fasta_fp.write('>%s\n' % right_template_name)
    right_template_fasta_fp.write('%s\n' % right_template_seq)
    right_template_fasta_fp.close()

    round1_paf_file = os.path.join(out_dir, 'round1.paf')
    round1_left_anchor_paf_file  = os.path.join(out_dir, 'round1.left.paf')
    round1_right_anchor_paf_file = os.path.join(out_dir, 'round1.right.paf')
    preset = tk.get_preset_for_minimap2(platform)

    cmd = f'{minimap2} -A 2 -B 6 -c --eqx -t {num_threads} -x {preset} {left_template_fasta_file} {in_fastq_file} > {round1_left_anchor_paf_file} 2> /dev/null'
    tk.run_system_cmd(cmd)

    cmd = f'{minimap2} -A 2 -B 6 -c --eqx -t {num_threads} -x {preset} {right_template_fasta_file} {in_fastq_file} > {round1_right_anchor_paf_file} 2> /dev/null'
    tk.run_system_cmd(cmd)

    cmd = f'cat {round1_left_anchor_paf_file} {round1_right_anchor_paf_file} | sort -k1 - > {round1_paf_file}'
    tk.run_system_cmd(cmd)

    tk.rm(round1_left_anchor_paf_file)
    tk.rm(round1_right_anchor_paf_file)
    initial_estimation = round1_estimation_from_paf(round1_paf_file, repeat1, repeat2, left_anchor_len, right_anchor_len)

    return initial_estimation

def round1_estimation_from_paf(round1_paf_file, repeat1, repeat2, left_anchor_len, right_anchor_len):

    initial_estimation = Round1Estimation()
    round1_paf_fp = open(round1_paf_file, 'r')
    read_paf_list = list()
    while 1:
        line = round1_paf_fp.readline()
        if not line: break
        line = line.strip()
        if not line: continue

        col_list = line.strip().split('\t')
        paf = tk.PAF(col_list)

        if len(read_paf_list) == 0 or paf.qname == read_paf_list[0].qname:
            read_paf_list.append(paf)
        else:
            round1_estimation_for1read(read_paf_list, repeat1, repeat2, left_anchor_len, right_anchor_len, initial_estimation)
            read_paf_list.clear()
            read_paf_list.append(paf)

    round1_paf_fp.close()

    round1_estimation_for1read(read_paf_list, repeat1, repeat2, left_anchor_len, right_anchor_len, initial_estimation)

    for readname in initial_estimation.bad_reads_set:
        initial_estimation.repeat1_count_range_dict.pop(readname, None)
        initial_estimation.repeat2_count_range_dict.pop(readname, None)

    return initial_estimation

def round1_estimation_for1read(read_paf_list, repeat1, repeat2, left_anchor_len, right_anchor_len, initial_estimation):

    left_boundary_pos  = left_anchor_len
    right_boundary_pos = right_anchor_len

    left_anchor_paf_list  = list()
    right_anchor_paf_list = list()

    for paf in read_paf_list:
        if paf.mapq < 30: continue
        if paf.is_primary == False: continue
        if paf.tname[0:5] == 'left_':
            if paf.tstart <= left_boundary_pos and paf.tend >= left_boundary_pos:
                left_anchor_paf_list.append(paf)
        elif paf.tname[0:5] == 'right':
            if paf.tstart <= right_boundary_pos and paf.tend >= right_boundary_pos:
                right_anchor_paf_list.append(paf)
        else:
            tk.eprint('ERROR! unknown template: %s' % (paf.tname))
            sys.exit(1)

    if len(left_anchor_paf_list) != 1 or len(right_anchor_paf_list) != 1 : return
    left_paf  = left_anchor_paf_list[0]
    right_paf  = right_anchor_paf_list[0]

    if left_paf.strand == right_paf.strand: return

    readname = left_paf.qname
    left_paf_max_repeat_size  = int((left_paf.tend - left_boundary_pos)/repeat1.repeat_unit_size) + 5
    left_paf_min_repeat_size  = tk.calculate_repeat_size_from_exact_match(left_paf.cigar, left_paf.tstart,   left_boundary_pos,  repeat1.repeat_unit_size)

    a = max(0, left_paf_min_repeat_size - 20)
    b = int(left_paf_min_repeat_size / 2.0)
    left_paf_min_repeat_size = min(a, b)

    initial_estimation.repeat1_count_range_dict[readname] = (left_paf_min_repeat_size, left_paf_max_repeat_size)

    right_paf_max_repeat_size = int((right_paf.tend - right_boundary_pos)/repeat2.repeat_unit_size) + 5
    right_paf_min_repeat_size = tk.calculate_repeat_size_from_exact_match(right_paf.cigar, right_paf.tstart, right_boundary_pos, repeat2.repeat_unit_size)
    a = max(0, right_paf_min_repeat_size - 20)
    b = int(right_paf_min_repeat_size / 2.0)
    right_paf_min_repeat_size = min(a, b)

    initial_estimation.repeat2_count_range_dict[readname] = (right_paf_min_repeat_size, right_paf_max_repeat_size)

    if left_paf.strand == '+':
        candidate_start = left_paf.qstart
        candidate_end = right_paf.qlen - right_paf.qstart
        if candidate_end - candidate_start <= 0: initial_estimation.bad_reads_set.add(readname)
    else:
        candidate_start = right_paf.qstart
        candidate_end = left_paf.qlen - left_paf.qstart
        if candidate_end - candidate_start <= 0: initial_estimation.bad_reads_set.add(readname)

    initial_estimation.potential_repeat_region_dict[readname] = (candidate_start, candidate_end)
    return

def fastq_file_to_dict(in_fastq_file):

    fastq_dict = dict()

    in_fastq_fp = tk.gzopen(in_fastq_file)
    while 1:
        line1 = in_fastq_fp.readline()
        line2 = in_fastq_fp.readline()
        line3 = in_fastq_fp.readline()
        line4 = in_fastq_fp.readline()

        if not line1: break
        if not line2: break
        if not line3: break
        if not line4: break

        readname = line1.strip().split()[0][1:]
        fastq_dict[readname] = line1 + line2 + line3 + line4

    in_fastq_fp.close()

    return fastq_dict

def jointly_split_alleles_using_gmm (ploidy, repeat1, repeat2, final_estimation, in_fastq_file, out_dir):

    in_fastq_prefix  = os.path.splitext(os.path.split(in_fastq_file)[1])[0]
    out_prefix = os.path.join(out_dir, '%s.JointGMM' % (in_fastq_prefix))
    out_fastq_prefix = os.path.join(out_dir, '%s.JointGMM.qc_passed' % (in_fastq_prefix))
    repeat_size_file = os.path.join(out_dir, '%s.repeat_size.txt' % in_fastq_prefix)

    read_repeat_joint_count_dict = dict()
    for readname in final_estimation.repeat1_count_dict:
        if readname not in final_estimation.repeat2_count_dict: continue
        size1 = final_estimation.repeat1_count_dict[readname]
        size2 = final_estimation.repeat2_count_dict[readname]
        read_repeat_joint_count_dict[readname] = (size1, size2)

    repeat_region_list = [repeat1, repeat2]

    if len(read_repeat_joint_count_dict) < ploidy:
        tk.eprint('WARNING: No enough reads! input fastq file is: %s\n' % in_fastq_file)
        readinfo_dict = dict()
        output_repeat_size_file(read_repeat_joint_count_dict, readinfo_dict, in_fastq_file, repeat_region_list, repeat_size_file, ploidy)
        return

    if ploidy < 1:
        tk.eprint('ploidy must be >= 1 !\n')
        sys.exit(1)

    proba_cutoff = 0.95
    cov_type = 'tied'

    min_count1, max_count1, min_count2, max_count2  = analysis_outlier_2d(read_repeat_joint_count_dict)

    readname_list = list() # readnames after removing outliers
    read_repeat_count_list = list()

    for readname in read_repeat_joint_count_dict:
        repeat_count1, repeat_count2 = read_repeat_joint_count_dict[readname]
        if repeat_count1 < min_count1 or repeat_count1 > max_count1: continue
        if repeat_count2 < min_count2 or repeat_count2 > max_count2: continue
        readname_list.append(readname)
        read_repeat_count_list.append(repeat_count1)
        read_repeat_count_list.append(repeat_count2)

    num_data_points = len(readname_list)
    read_repeat_count_ndarray = np.array(read_repeat_count_list)
    read_repeat_count_ndarray = read_repeat_count_ndarray.reshape(num_data_points, 2)

    best_n_components = chose_best_num_components (read_repeat_count_ndarray, ploidy, cov_type)
    tk.eprint('NOTICE: number of alleles = %d' % best_n_components)
    final_gmm = GaussianMixture(n_components = best_n_components, covariance_type=cov_type, n_init = 100).fit(read_repeat_count_ndarray)
    old_read_label_list = list(final_gmm.predict(read_repeat_count_ndarray))
    proba2darray = final_gmm.predict_proba(read_repeat_count_ndarray)

    repeat_region_list[0].is_main = 1
    main_repeat_idx = 0
    for i in range(0, len(repeat_region_list)):
        if repeat_region_list[i].is_main:
            main_repeat_idx = i
            break

    old_cluster_mean_list = list()
    for means in final_gmm.means_:
        old_cluster_mean_list.append(means[main_repeat_idx])

    read_label_list, old_label_to_new_label_dict, new_label_to_old_label_dict = sort_label_by_cluster_mean(old_read_label_list, old_cluster_mean_list)

    readinfo_dict = dict()
    for i in range(0, len(readname_list)):
        readname = readname_list[i]
        repeat_count1, repeat_count2 = read_repeat_joint_count_dict[readname]
        read_label = read_label_list[i]
        max_prob = max(proba2darray[i])
        readinfo = Readinfo(readname)
        readinfo.label = read_label
        readinfo.max_prob = max_prob
        readinfo.proba_list = list(proba2darray[i])
        readinfo.repeat_size1 = repeat_count1
        readinfo.repeat_size2 = repeat_count2
        readinfo_dict[readname] = readinfo

    each_allele_repeat_count_3d_list = [0] * best_n_components # each_allele_repeat_count_3d_list[label][repeat_id] = list of repeat size
    for label in range(0, len(each_allele_repeat_count_3d_list)):
        each_allele_repeat_count_3d_list[label] = [0] * 2
        for repeat_id in range(0, 2):
            each_allele_repeat_count_3d_list[label][repeat_id] = list()

    for readname in readinfo_dict:
        readinfo = readinfo_dict[readname]
        each_allele_repeat_count_3d_list[readinfo.label][0].append(readinfo.repeat_size1)
        each_allele_repeat_count_3d_list[readinfo.label][1].append(readinfo.repeat_size2)

    allele_predicted_repeat_size_2d_list = [0] * best_n_components
    for label in range(0, len(allele_predicted_repeat_size_2d_list)):
        allele_predicted_repeat_size_2d_list[label] = [0] * 2
        allele_predicted_repeat_size_2d_list[label][0] = int(np.median(each_allele_repeat_count_3d_list[label][0]) + 0.5)
        allele_predicted_repeat_size_2d_list[label][1] = int(np.median(each_allele_repeat_count_3d_list[label][1]) + 0.5)

    score_cut_off = calculate_log_likelyhood_cutoff(final_gmm, 0.99)
    label_qc_failed_reads(readinfo_dict, final_gmm, proba_cutoff, score_cut_off)

    tk.eprint('NOTICE: writing to repeat size file.')
    output_repeat_size_file(read_repeat_joint_count_dict, readinfo_dict, in_fastq_file, repeat_region_list, repeat_size_file, ploidy)
    tk.eprint('NOTICE: writing to output fastq files.')
    joint_gmm_output_fastq(in_fastq_file, readinfo_dict, best_n_components, out_fastq_prefix)

    tk.eprint('NOTICE: writing to output summary file.')
    joint_gmm_output_summary_file(final_gmm, in_fastq_file, readinfo_dict, allele_predicted_repeat_size_2d_list, repeat_region_list, out_prefix)

    tk.eprint('NOTICE: plotting figures.')
    joint_gmm_plot_repeat_counts(readinfo_dict, allele_predicted_repeat_size_2d_list, repeat_region_list, out_prefix)

    joint_gmm_scatter_plot_with_contour (read_repeat_joint_count_dict, final_gmm, score_cut_off, repeat_region_list, out_prefix)

    return

def joint_gmm_scatter_plot_with_contour (read_repeat_joint_count_dict, final_gmm, score_cut_off, repeat_region_list, out_prefix):

    scatter_plot_file = out_prefix + '.scatter.png'

    xlabel = '%s repeat size' % (repeat_region_list[0].repeat_id)
    ylabel = '%s repeat size' % (repeat_region_list[1].repeat_id)

    X = list()
    Y = list()

    for readname in read_repeat_joint_count_dict:
        x, y = read_repeat_joint_count_dict[readname]
        X.append(x)
        Y.append(y)

    X = np.array(X)
    Y = np.array(Y)

    X, Y, Z= countxy(X, Y)
    fig, ax = plt.subplots()
    fig.set_size_inches(6, 4)

    ax.scatter(X, Y, c=Z, s=15, edgecolor='')

    plt.xlabel(xlabel)
    plt.ylabel(ylabel)

    norm = Normalize(vmin = np.min(Z), vmax = np.max(Z))
    cbar = fig.colorbar(cm.ScalarMappable(norm = norm), ax=ax)
    cbar.ax.set_ylabel('Count')

    xmin = min(X)
    xmax = max(X)
    ymin = min(Y)
    ymax = max(Y)

    a = np.linspace(xmin, xmax, 200)
    b = np.linspace(ymin, ymax, 200)

    A, B = np.meshgrid(a, b)
    AA = np.array([A.ravel(), B.ravel()]).T
    C = final_gmm.score_samples(AA)
    C = C.reshape(A.shape)

    CS = plt.contour(A, B, C, levels=[score_cut_off], linestyles = 'dashed', colors = 'grey')

    plt.savefig(scatter_plot_file, dpi = 300)
    plt.close('all')
    return

def countxy(x, y):
    count_dict = dict()
    for i in range(0, len(x)):
        key = '%d\t%d' % (x[i], y[i])
        if key not in count_dict:
            count_dict[key] = 1
        else:
            count_dict[key] += 1
    x_list = list()
    y_list = list()
    z_list = list()
    for key in count_dict:
        x, y = key.split('\t')
        x = int(x)
        y = int(y)
        z = count_dict[key]
        x_list.append(x)
        y_list.append(y)
        z_list.append(z)

    return x_list, y_list, z_list

def calculate_log_likelyhood_cutoff(gmm, ci):
    X, labels = gmm.sample(100000)
    score_list = gmm.score_samples(X)
    X_score_list = list()
    for i in range(0, len(X)):
        X_score = [X[i], score_list[i]]
        X_score_list.append(X_score)

    X_score_list.sort(key = lambda x:x[1], reverse = True)
    cut_off_idx = int(len(X_score_list) * float(ci))
    log_likelyhood_cutoff = X_score_list[cut_off_idx][1]

    return log_likelyhood_cutoff

def joint_gmm_plot_repeat_counts(readinfo_dict, allele_predicted_repeat_size_2d_list, repeat_region_list, out_prefix):

    hist2d_figure_file = out_prefix + '.hist2d.png'
    repeat1_hist_figure_file  = out_prefix + '.%s.hist.png' % (repeat_region_list[0].repeat_id)
    repeat2_hist_figure_file  = out_prefix + '.%s.hist.png' % (repeat_region_list[1].repeat_id)

    num_alleles = len(allele_predicted_repeat_size_2d_list)
    x_list = list()
    y_list = list()

    x_2d_list = [0] * num_alleles
    y_2d_list = [0] * num_alleles
    for i in range(0, num_alleles):
        x_2d_list[i] = list()
        y_2d_list[i] = list()
    for readname in readinfo_dict:
        readinfo = readinfo_dict[readname]
        if readinfo.qc_passed == 0: continue

        x = readinfo.repeat_size1
        y = readinfo.repeat_size2
        x_list.append(x)
        y_list.append(y)

        x_2d_list[readinfo.label].append(x)
        y_2d_list[readinfo.label].append(y)

    xmin = int(min(x_list))
    xmax = int(max(x_list))+1
    ymin = int(min(y_list))
    ymax = int(max(y_list))+1

    if xmax - xmin <= 200:
        b1 = range(xmin - 1, xmax + 2)
    else:
        b1 = range(xmin - 1, xmax + 2, int(float(xmax - xmin)/200.0 + 0.5))

    if ymax - ymin <= 200:
        b2 = range(ymin - 1, ymax + 2)
    else:
        b2 = range(ymin - 1, ymax + 2, int(float(ymax - ymin)/200.0 + 0.5))

    repeat1_predicted_size_list = list()
    repeat2_predicted_size_list = list()
    for label in range(0, len(allele_predicted_repeat_size_2d_list)):
        repeat1_predicted_size_list.append(allele_predicted_repeat_size_2d_list[label][0])
        repeat2_predicted_size_list.append(allele_predicted_repeat_size_2d_list[label][1])

    plot_hist2d(x_list, y_list, b1, b2, repeat_region_list[0].repeat_id, repeat_region_list[1].repeat_id, hist2d_figure_file)
    plot_hist1d(x_2d_list, b1, repeat_region_list[0].repeat_id, repeat1_predicted_size_list, repeat1_hist_figure_file)
    plot_hist1d(y_2d_list, b2, repeat_region_list[1].repeat_id, repeat2_predicted_size_list, repeat2_hist_figure_file)

    return

def plot_hist1d(x_2d_list, b, repeat_id, predicted_size_list, out_file):

    plt.figure (figsize=(6, 4))

    for x in x_2d_list:
        plt.hist(x, bins = b)

    for repeat_size in predicted_size_list:
        plt.axvline(x=repeat_size+0.5, color = 'grey', linestyle = ':')

    plt.title('Repeat size distribution (%s)' % repeat_id)
    plt.xlabel('repeat size')
    plt.ylabel('number of reads')
    if debug:
        plt.xlim(0, 150)

    plt.savefig(out_file, dpi=300)
    plt.close('all')

    return

def plot_hist2d(x_list, y_list, b1, b2, repeat_id1, repeat_id2, out_file):

    plt.figure (figsize=(8, 4))
    plt.hist2d (x_list, y_list, [b1, b2], cmap = 'binary')
    plt.colorbar()
    plt.title('2D histogram of repeat size')
    plt.xlabel('repeat size (%s)' % repeat_id1)
    plt.ylabel('repeat size (%s)' % repeat_id2)
    xmin = min(x_list)
    xmax = max(x_list)

    if debug:
        plt.xlim(xmin, xmax)
        plt.ylim(0, (xmax-xmin)/2)

    plt.savefig(out_file, dpi=300)
    plt.close('all')

    return

def joint_gmm_output_summary_file(final_gmm, in_fastq_file, readinfo_dict, allele_predicted_repeat_size_2d_list, repeat_region_list, out_prefix):

    out_summray_file = out_prefix + '.summary.txt'
    out_summray_fp = open(out_summray_file, 'w')
    summary_header = '#input_fastq\tmethod'
    summary_info = '%s\tJointGMM' % (in_fastq_file)

    # allele_predicted_repeat_size_2d_list[label][repeat_id] = predicted_repeat_size
    num_alleles = len(allele_predicted_repeat_size_2d_list)

    cov1 = final_gmm.covariances_[0][0]
    cov2 = final_gmm.covariances_[1][1]
    summary_header += '\tnum_alleles\tgmm_cov_repeat1\tgmm_cov_repeat2'
    summary_info   += '\t%d\t%.4f\t%.4f' % (num_alleles, cov1, cov2)

    allele_num_reads_list = [0] * num_alleles
    for readname in readinfo_dict:
        readinfo = readinfo_dict[readname]
        if readinfo.qc_passed == 0: continue
        allele_num_reads_list[readinfo.label] += 1

    for label in range(0, num_alleles):
        allele_id = label + 1
        summary_header += '\tallele%d_num_reads' % (allele_id)
        summary_info   += '\t%d' % allele_num_reads_list[label]
        for i in range(0, 2):
            summary_header += '\t%s_repeat_size%d' % (repeat_region_list[i].repeat_id, allele_id)
            summary_info   += '\t%d' % (allele_predicted_repeat_size_2d_list[label][i])

    out_summray_fp.write(summary_header + '\n')
    out_summray_fp.write(summary_info + '\n')

    return

def joint_gmm_output_fastq(in_fastq_file, readinfo_dict, num_alleles, out_prefix):

    out_allele_fastq_file_list = list()
    for label in range(0, num_alleles):
        allele_id = label + 1
        out_allele_fastq_file = out_prefix + '.allele%d.fastq' % (allele_id)
        out_allele_fastq_file_list.append(out_allele_fastq_file)

    out_allele_fastq_fp_list = list()
    for i in range(0, len(out_allele_fastq_file_list)):
        out_allele_fastq_fp = open(out_allele_fastq_file_list[i], 'w')
        out_allele_fastq_fp_list.append(out_allele_fastq_fp)

    if '.gz' == in_fastq_file[-3:]:
        in_fastq_fp = gzip.open(in_fastq_file, 'rt')
    else:
        in_fastq_fp = open(in_fastq_file, 'rt')

    while 1:
        line1 = in_fastq_fp.readline()
        line2 = in_fastq_fp.readline()
        line3 = in_fastq_fp.readline()
        line4 = in_fastq_fp.readline()

        if not line1: break
        if not line2: break
        if not line3: break
        if not line4: break

        readname = line1.strip().split()[0][1:]
        if readname not in readinfo_dict: continue
        if readinfo_dict[readname].qc_passed == 0: continue

        label = readinfo_dict[readname].label
        out_allele_fastq_fp_list[label].write(line1 + line2 + line3 + line4)

    in_fastq_fp.close()

    for i in range(0, len(out_allele_fastq_fp_list)):
        out_allele_fastq_fp_list[i].close()

    return

def label_qc_failed_reads(readinfo_dict, final_gmm, proba_cutoff, score_cut_off):

    for readname in readinfo_dict:
        readinfo = readinfo_dict[readname]
        qc_failed = 0

        x = [readinfo.repeat_size1, readinfo.repeat_size2]
        x = np.array(x).reshape(1, 2)
        score = final_gmm.score_samples(x)
        if score[0] < score_cut_off: qc_failed = 1

        if readinfo.max_prob < proba_cutoff: qc_failed = 1

        if qc_failed: readinfo_dict[readname].qc_passed = 0

    return

def bic_best_num_components (X, max_num_components, cov_type):

    bic_list = list()
    bic_list.append(0)

    fold = 0
    if len(X) < 1000:
        fold = int(1000 / len(X))

    Y = np.concatenate((X, X))
    for j in range(0, fold):
        Y = np.concatenate((Y, X))

    error = np.random.rand(Y.shape[0], Y.shape[1]) - 0.5
    Y = Y + error
    for n in range(1, max_num_components+1):
        gmm = GaussianMixture(n_components=n, covariance_type=cov_type, n_init=20).fit(Y)
        bic = gmm.bic(Y)
        bic_list.append(bic)

    min_bic = 1e99
    min_bic_n_components = 0
    for i in range(1, len(bic_list)):
        if bic_list[i] < min_bic:
            min_bic = bic_list[i]
            min_bic_n_components = i

    return min_bic_n_components

def chose_best_num_components (X, max_num_components, cov_type):

    min_mean_distance = 3
    bic_num_component = bic_best_num_components (X, max_num_components, cov_type)
    best_n_components = bic_num_component
    for n in range(bic_num_component, 0, -1):
        gmm = GaussianMixture(n_components=n, covariance_type=cov_type, n_init=100).fit(X)
        cov1 = gmm.covariances_[0][0]
        cov2 = gmm.covariances_[1][1]
        too_close = 0
        for i in range(0, n):
            for j in range(i+1, n):
                mean1 = gmm.means_[i]
                mean2 = gmm.means_[j]
                if abs(mean1[0]-mean2[0]) < min_mean_distance and abs(mean1[1]-mean2[1]) < min_mean_distance:
                    too_close = 1
                    break
                if abs(mean1[0]-mean2[0]) < cov1 * 2 and abs(mean1[1]-mean2[1]) < cov2 * 2:
                    too_close = 1
                    break

        if too_close == 0:
            best_n_components = n
            break

    return best_n_components

def analysis_outlier(read_repeat_count_dict):

    read_repeat_count_list = list()
    for readname in read_repeat_count_dict:
        repeat_count = read_repeat_count_dict[readname]
        read_repeat_count_list.append(repeat_count)

    min_repeat_count_cutoff, max_repeat_count_cutoff = get_outlier_cutoff_from_list(read_repeat_count_list)

    return min_repeat_count_cutoff, max_repeat_count_cutoff

def analysis_outlier_2d(read_repeat_joint_count_dict):

    repeat_count1_list = list()
    repeat_count2_list = list()

    for readname in read_repeat_joint_count_dict:
        repeat_count1, repeat_count2 = read_repeat_joint_count_dict[readname]
        repeat_count1_list.append(repeat_count1)
        repeat_count2_list.append(repeat_count2)

    min_count1, max_count1 = get_outlier_cutoff_from_list(repeat_count1_list)
    min_count2, max_count2 = get_outlier_cutoff_from_list(repeat_count2_list)

    return min_count1, max_count1, min_count2, max_count2

def get_outlier_cutoff_from_list(repeat_count_list):

    mean = np.mean(repeat_count_list)
    std = np.std(repeat_count_list)

    min_repeat_count_cutoff = mean - 3 * std
    if min_repeat_count_cutoff < 0: min_repeat_count_cutoff = 0
    max_repeat_count_cutoff = mean + 3 * std

    return min_repeat_count_cutoff, max_repeat_count_cutoff

def sort_label_by_cluster_mean(old_read_label_list, cluster_mean_list):

    l = list()
    for i in range(0, len(cluster_mean_list)):
        label = i
        cluster_mean = cluster_mean_list[i]
        l.append((label, cluster_mean))

    l = sorted(l, key=lambda x:x[1])

    old_label_to_new_label_dict = dict()
    new_label_to_old_label_dict = dict()

    for i in range(0, len(l)):
        new_label = i
        old_label = l[i][0]
        old_label_to_new_label_dict[old_label] = new_label
        new_label_to_old_label_dict[new_label] = old_label

    new_read_label_list = list()

    for i in range(0, len(old_read_label_list)):
        old_label = old_read_label_list[i]
        new_label = old_label_to_new_label_dict[old_label]
        new_read_label_list.append(new_label)

    return new_read_label_list, old_label_to_new_label_dict, new_label_to_old_label_dict

class Repeat:
    def __init__(self):
        self.repeat_id = ''
        self.chrom = ''
        self.start = -1
        self.end = -1
        self.repeat_unit = ''
        self.repeat_unit_size = 0
        self.min_size = 0
        self.max_size = 1000
        self.round1_min_size = 0
        self.round1_max_size = 1000
        self.is_main = 0

    def init_from_string(self, string):
        col_list = string.split(':')
        if len(col_list) != 5:
            tk.eprint('ERROR! --repeat1 and --repeat2 should be in this format: chr:start:end:repeat_unit:max_size (e.g. chr4:3074876:3074933:CAG:200')
            sys.exit(1)

        self.chrom, self.start, self.end, self.repeat_unit, self.max_size = col_list
        self.start = int(self.start)
        self.end   = int(self.end)
        self.repeat_unit_size = len(self.repeat_unit)
        self.min_size = 0
        self.max_size = int(self.max_size)
        self.repeat_id = '_'.join(col_list[0:4])
        self.start -= 1
        return self

class Round1Estimation:
    def __init__(self):
        self.repeat1_count_range_dict = dict()
        self.repeat2_count_range_dict = dict()
        self.potential_repeat_region_dict = dict()
        self.bad_reads_set = set()

    def output_repeat2_boundaries(self):
        out_string = ''
        for readname in self.repeat2_count_range_dict:
            lower_bound, upper_bound = self.repeat2_count_range_dict[readname]
            out_string += f'{readname}\t{lower_bound}\t{upper_bound}\n'

        return out_string

class RepeatSize:
    def __init__(self):
        self.repeat1_count_dict = dict()
        self.repeat2_count_dict = dict()
        self.step_size1 = 1
        self.step_size2 = 1

if __name__ == '__main__':
    main()


import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
from torchvision import ops
# ops.DeformConv2d()

from .correlation_package.correlation import Correlation

def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1, activation=True):
    if activation:
        return nn.Sequential(
            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
                      padding=padding, dilation=dilation, bias=True),
            nn.LeakyReLU(0.1))
    else:
        return nn.Sequential(
            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
                      padding=padding, dilation=dilation, bias=True))

def predict_flow(in_planes):
    return nn.Conv2d(in_planes, 2, kernel_size=3, stride=1, padding=1, bias=True)

def predict_mask(in_planes):
    return nn.Conv2d(in_planes, 1, kernel_size=3, stride=1, padding=1, bias=True)

def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
    return nn.ConvTranspose2d(in_planes, out_planes, kernel_size, stride, padding, bias=True)

def deformable_conv(in_planes, out_planes, kernel_size=3, strides=1, padding=1, use_bias=True):
    return ops.DeformConv2d(in_planes, out_planes, kernel_size, strides, padding, bias=use_bias)

def upsample_kernel2d(w, device):
    c = w // 2
    kernel = 1 - torch.abs(c - torch.arange(w, dtype=torch.float32, device=device)) / (c + 1)
    kernel = kernel.repeat(w).view(w,-1) * kernel.unsqueeze(1)
    return kernel.view(1, 1, w, w)

def downsample_kernel2d(w, device):
    kernel = ((w + 1) - torch.abs(w - torch.arange(w * 2 + 1, dtype=torch.float32, device=device))) / (2 * w + 1)
    kernel = kernel.repeat(w).view(w,-1) * kernel.unsqueeze(1)
    return kernel.view(1, 1, w * 2 + 1, w * 2 + 1)

def Upsample(img, factor):
    if factor == 1:
        return img
    B, C, H, W = img.shape
    batch_img = img.view(B*C, 1, H, W)
    batch_img = F.pad(batch_img, [0, 1, 0, 1], mode='replicate')
    kernel = upsample_kernel2d(factor * 2 - 1, img.device)
    upsamp_img = F.conv_transpose2d(batch_img, kernel, stride=factor, padding=(factor-1))
    upsamp_img = upsamp_img[:, :, : -1, :-1]
    _, _, H_up, W_up = upsamp_img.shape
    return upsamp_img.view(B, C, H_up, W_up)

def Downsample(img, factor):
    if factor == 1:
        return img
    B, C, H, W = img.shape
    batch_img = img.view(B*C, 1, H, W)
    kernel = downsample_kernel2d(factor // 2, img.device)
    upsamp_img = F.conv2d(batch_img, kernel, stride=factor, padding=factor//2)
    upsamp_nom = F.conv2d(torch.ones_like(batch_img), kernel, stride=factor, padding=factor//2)
    _, _, H_up, W_up = upsamp_img.shape
    upsamp_img = upsamp_img.view(B, C, H_up, W_up)
    upsamp_nom = upsamp_nom.view(B, C, H_up, W_up)
    return upsamp_img / upsamp_nom

class MaskFlownet_S(nn.Module):
    """
    PWC-DC net. add dilation convolution and densenet connections
    """
    def __init__(self, config = None, **kwargs):
        """
        input: md --- maximum displacement (for correlation. default: 4), after warpping
        """
        super(MaskFlownet_S, self).__init__()
        self.scale = 20. * config.network.flow_multiplier.get(1.)
        md = 4
        self.md = md
        self.strides = [64, 32, 16, 8, 4]
        self.deform_bias = config.network.deform_bias.get(True)
        self.upfeat_ch = config.network.upfeat_ch.get([16, 16, 16, 16])

        self.conv1a  = conv(3,   16, kernel_size=3, stride=2)
        self.conv1b = conv(16, 16, kernel_size=3, stride=1)
        self.conv1c  = conv(16, 16, kernel_size=3, stride=1)
        self.conv2a  = conv(16,  32, kernel_size=3, stride=2)
        self.conv2b = conv(32, 32, kernel_size=3, stride=1)
        self.conv2c  = conv(32, 32, kernel_size=3, stride=1)
        self.conv3a  = conv(32,  64, kernel_size=3, stride=2)
        self.conv3b = conv(64, 64, kernel_size=3, stride=1)
        self.conv3c  = conv(64, 64, kernel_size=3, stride=1)
        self.conv4a  = conv(64,  96, kernel_size=3, stride=2)
        self.conv4b = conv(96, 96, kernel_size=3, stride=1)
        self.conv4c  = conv(96, 96, kernel_size=3, stride=1)
        self.conv5a  = conv(96, 128, kernel_size=3, stride=2)
        self.conv5b = conv(128, 128, kernel_size=3, stride=1)
        self.conv5c  = conv(128, 128, kernel_size=3, stride=1)
        self.conv6a = conv(128, 196, kernel_size=3, stride=2)
        self.conv6b  = conv(196, 196, kernel_size=3, stride=1)
        self.conv6c  = conv(196, 196, kernel_size=3, stride=1)

        self.corr    = Correlation(pad_size=md, kernel_size=1, max_displacement=md, stride1=1, stride2=1, corr_multiply=1)
        self.leakyRELU = nn.LeakyReLU(0.1)

        nd = (2*md+1)**2
        dd = np.cumsum([128,128,96,64,32])

        od = nd
        self.conv6_0 = conv(od,      128, kernel_size=3, stride=1)
        self.conv6_1 = conv(od+dd[0],128, kernel_size=3, stride=1)
        self.conv6_2 = conv(od+dd[1],96,  kernel_size=3, stride=1)
        self.conv6_3 = conv(od+dd[2],64,  kernel_size=3, stride=1)
        self.conv6_4 = conv(od+dd[3],32,  kernel_size=3, stride=1)
        self.pred_flow6 = predict_flow(od + dd[4])
        self.pred_mask6 = predict_mask(od + dd[4])
        self.upfeat5 = deconv(od+dd[4], self.upfeat_ch[0], kernel_size=4, stride=2, padding=1)
        # self.deconv6 = deconv(2, 2, kernel_size=4, stride=2, padding=1)
        # self.upfeat6 = deconv(od+dd[4], 2, kernel_size=4, stride=2, padding=1)

        # od = nd+128+4
        od = nd+128+18
        self.conv5_0 = conv(od,      128, kernel_size=3, stride=1)
        self.conv5_1 = conv(od+dd[0],128, kernel_size=3, stride=1)
        self.conv5_2 = conv(od+dd[1],96,  kernel_size=3, stride=1)
        self.conv5_3 = conv(od+dd[2],64,  kernel_size=3, stride=1)
        self.conv5_4 = conv(od+dd[3],32,  kernel_size=3, stride=1)
        self.pred_flow5 = predict_flow(od + dd[4])
        self.pred_mask5 = predict_mask(od + dd[4])
        self.upfeat4 = deconv(od+dd[4], self.upfeat_ch[1], kernel_size=4, stride=2, padding=1)
        # self.deconv5 = deconv(2, 2, kernel_size=4, stride=2, padding=1)
        # self.upfeat5 = deconv(od+dd[4], 2, kernel_size=4, stride=2, padding=1)

        # od = nd+96+4
        od = nd+96+18
        self.conv4_0 = conv(od,      128, kernel_size=3, stride=1)
        self.conv4_1 = conv(od+dd[0],128, kernel_size=3, stride=1)
        self.conv4_2 = conv(od+dd[1],96,  kernel_size=3, stride=1)
        self.conv4_3 = conv(od+dd[2],64,  kernel_size=3, stride=1)
        self.conv4_4 = conv(od+dd[3],32,  kernel_size=3, stride=1)
        self.pred_flow4 = predict_flow(od + dd[4])
        self.pred_mask4 = predict_mask(od + dd[4])
        self.upfeat3 = deconv(od+dd[4], self.upfeat_ch[2], kernel_size=4, stride=2, padding=1)
        # self.deconv4 = deconv(2, 2, kernel_size=4, stride=2, padding=1)
        # self.upfeat4 = deconv(od+dd[4], 2, kernel_size=4, stride=2, padding=1)

        # od = nd+64+4
        od = nd+64+18
        self.conv3_0 = conv(od,      128, kernel_size=3, stride=1)
        self.conv3_1 = conv(od+dd[0],128, kernel_size=3, stride=1)
        self.conv3_2 = conv(od+dd[1],96,  kernel_size=3, stride=1)
        self.conv3_3 = conv(od+dd[2],64,  kernel_size=3, stride=1)
        self.conv3_4 = conv(od+dd[3],32,  kernel_size=3, stride=1)
        self.pred_flow3 = predict_flow(od + dd[4])
        self.pred_mask3 = predict_mask(od + dd[4])
        self.upfeat2 = deconv(od+dd[4], self.upfeat_ch[3], kernel_size=4, stride=2, padding=1)
        # self.deconv3 = deconv(2, 2, kernel_size=4, stride=2, padding=1)
        # self.upfeat3 = deconv(od+dd[4], 2, kernel_size=4, stride=2, padding=1)

        # od = nd+32+4
        od = nd+32+18
        self.conv2_0 = conv(od,      128, kernel_size=3, stride=1)
        self.conv2_1 = conv(od+dd[0],128, kernel_size=3, stride=1)
        self.conv2_2 = conv(od+dd[1],96,  kernel_size=3, stride=1)
        self.conv2_3 = conv(od+dd[2],64,  kernel_size=3, stride=1)
        self.conv2_4 = conv(od+dd[3],32,  kernel_size=3, stride=1)
        self.pred_flow2 = predict_flow(od + dd[4])

        self.dc_conv1 = conv(od+dd[4], 128, kernel_size=3, stride=1, padding=1,  dilation=1)
        self.dc_conv2 = conv(128,      128, kernel_size=3, stride=1, padding=2,  dilation=2)
        self.dc_conv3 = conv(128,      128, kernel_size=3, stride=1, padding=4,  dilation=4)
        self.dc_conv4 = conv(128,      96,  kernel_size=3, stride=1, padding=8,  dilation=8)
        self.dc_conv5 = conv(96,       64,  kernel_size=3, stride=1, padding=16, dilation=16)
        self.dc_conv6 = conv(64,       32,  kernel_size=3, stride=1, padding=1,  dilation=1)
        self.dc_conv7 = predict_flow(32)

        # self.upfeat5 = deconv()

        self.deform5 = deformable_conv(128, 128)
        self.deform4 = deformable_conv(96, 96)
        self.deform3 = deformable_conv(64, 64)
        self.deform2 = deformable_conv(32, 32)

        self.conv5f = conv(16, 128, kernel_size=3, stride=1, padding=1, activation=False)
        self.conv4f = conv(16, 96, kernel_size=3, stride=1, padding=1, activation=False)
        self.conv3f = conv(16, 64, kernel_size=3, stride=1, padding=1, activation=False)
        self.conv2f = conv(16, 32, kernel_size=3, stride=1, padding=1, activation=False)

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                nn.init.kaiming_normal(m.weight.data, mode='fan_in')
                if m.bias is not None:
                    m.bias.data.zero_()

    def warp(self, x, flo):
        """
        warp an image/tensor (im2) back to im1, according to the optical flow
        x: [B, C, H, W] (im2)
        flo: [B, 2, H, W] flow
        """
        B, C, H, W = x.size()
        # mesh grid
        xx = torch.arange(0, W).view(1,-1).repeat(H,1)
        yy = torch.arange(0, H).view(-1,1).repeat(1,W)
        xx = xx.view(1,1,H,W).repeat(B,1,1,1)
        yy = yy.view(1,1,H,W).repeat(B,1,1,1)
        grid = torch.cat((xx,yy),1).float()

        device = x.device
        grid = grid.to(device)
        # vgrid = Variable(grid) + flo
        vgrid = Variable(grid) + torch.flip(flo, [1])

        # scale grid to [-1,1]
        vgrid[:,0,:,:] = 2.0*vgrid[:,0,:,:].clone() / max(W-1,1)-1.0
        vgrid[:,1,:,:] = 2.0*vgrid[:,1,:,:].clone() / max(H-1,1)-1.0

        vgrid = vgrid.permute(0,2,3,1)
        # vgrid = vgrid.permute(0,2,3,1).clamp(-1.1, 1.1)
        output = nn.functional.grid_sample(x, vgrid, align_corners=True)
        mask = torch.autograd.Variable(torch.ones(x.size())).to(device)
        mask = nn.functional.grid_sample(mask, vgrid, align_corners=True)

        # if W==128:
        # np.save('mask.npy', mask.cpu().data.numpy())
        # np.save('warp.npy', output.cpu().data.numpy())

        mask[mask<0.9999] = 0
        mask[mask>0] = 1

        return output*mask

    def forward(self, im1, im2):
        # im1 = x[:,:3,:,:]
        # im2 = x[:,3:,:,:]

        c11 = self.conv1c(self.conv1b(self.conv1a(im1)))
        c21 = self.conv1c(self.conv1b(self.conv1a(im2)))
        c12 = self.conv2c(self.conv2b(self.conv2a(c11)))
        c22 = self.conv2c(self.conv2b(self.conv2a(c21)))
        c13 = self.conv3c(self.conv3b(self.conv3a(c12)))
        c23 = self.conv3c(self.conv3b(self.conv3a(c22)))
        c14 = self.conv4c(self.conv4b(self.conv4a(c13)))
        c24 = self.conv4c(self.conv4b(self.conv4a(c23)))
        c15 = self.conv5c(self.conv5b(self.conv5a(c14)))
        c25 = self.conv5c(self.conv5b(self.conv5a(c24)))
        c16 = self.conv6c(self.conv6b(self.conv6a(c15)))
        c26 = self.conv6c(self.conv6b(self.conv6a(c25)))

        corr6 = self.corr(c16, c26)
        corr6 = self.leakyRELU(corr6)

        x = torch.cat((self.conv6_0(corr6), corr6),1)
        x = torch.cat((self.conv6_1(x), x),1)
        x = torch.cat((self.conv6_2(x), x),1)
        x = torch.cat((self.conv6_3(x), x),1)
        x = torch.cat((self.conv6_4(x), x),1)
        flow6 = self.pred_flow6(x)
        mask6 = self.pred_mask6(x)

        feat5 = self.leakyRELU(self.upfeat5(x))
        flow5 = Upsample(flow6, 2)
        mask5 = Upsample(mask6, 2)
        warp5 = (flow5*self.scale/self.strides[1]).unsqueeze(1)
        warp5 = torch.repeat_interleave(warp5, 9, 1)
        S1, S2, S3, S4, S5 = warp5.shape
        warp5 = warp5.view(S1, S2*S3, S4, S5)
        warp5 = self.deform5(c25, warp5)
        tradeoff5 = feat5
        warp5 = (warp5 * F.sigmoid(mask5)) + self.conv5f(tradeoff5)
        warp5 = self.leakyRELU(warp5)
        corr5 = self.corr(c15, warp5)
        corr5 = self.leakyRELU(corr5)
        x = torch.cat((corr5, c15, feat5, flow5), 1)
        x = torch.cat((self.conv5_0(x), x),1)
        x = torch.cat((self.conv5_1(x), x),1)
        x = torch.cat((self.conv5_2(x), x),1)
        x = torch.cat((self.conv5_3(x), x),1)
        x = torch.cat((self.conv5_4(x), x),1)
        flow5 = flow5 + self.pred_flow5(x)
        mask5 = self.pred_mask5(x)

        feat4 = self.leakyRELU(self.upfeat4(x))
        flow4 = Upsample(flow5, 2)
        mask4 = Upsample(mask5, 2)
        warp4 = (flow4*self.scale/self.strides[2]).unsqueeze(1)
        warp4 = torch.repeat_interleave(warp4, 9, 1)
        S1, S2, S3, S4, S5 = warp4.shape
        warp4 = warp4.view(S1, S2*S3, S4, S5)
        warp4 = self.deform4(c24, warp4)
        tradeoff4 = feat4
        warp4 = (warp4 * F.sigmoid(mask4)) + self.conv4f(tradeoff4)
        warp4 = self.leakyRELU(warp4)
        corr4 = self.corr(c14, warp4)
        corr4 = self.leakyRELU(corr4)
        x = torch.cat((corr4, c14, feat4, flow4), 1)
        x = torch.cat((self.conv4_0(x), x),1)
        x = torch.cat((self.conv4_1(x), x),1)
        x = torch.cat((self.conv4_2(x), x),1)
        x = torch.cat((self.conv4_3(x), x),1)
        x = torch.cat((self.conv4_4(x), x),1)
        flow4 = flow4 + self.pred_flow4(x)
        mask4 = self.pred_mask4(x)

        feat3 = self.leakyRELU(self.upfeat3(x))
        flow3 = Upsample(flow4, 2)
        mask3 = Upsample(mask4, 2)
        warp3 = (flow3*self.scale/self.strides[3]).unsqueeze(1)
        warp3 = torch.repeat_interleave(warp3, 9, 1)
        S1, S2, S3, S4, S5 = warp3.shape
        warp3 = warp3.view(S1, S2*S3, S4, S5)
        warp3 = self.deform3(c23, warp3)
        tradeoff3 = feat3
        warp3 = (warp3 * F.sigmoid(mask3)) + self.conv3f(tradeoff3)
        warp3 = self.leakyRELU(warp3)
        corr3 = self.corr(c13, warp3)
        corr3 = self.leakyRELU(corr3)
        x = torch.cat((corr3, c13, feat3, flow3), 1)
        x = torch.cat((self.conv3_0(x), x),1)
        x = torch.cat((self.conv3_1(x), x),1)
        x = torch.cat((self.conv3_2(x), x),1)
        x = torch.cat((self.conv3_3(x), x),1)
        x = torch.cat((self.conv3_4(x), x),1)
        flow3 = flow3 + self.pred_flow3(x)
        mask3 = self.pred_mask3(x)

        feat2 = self.leakyRELU(self.upfeat2(x))
        flow2 = Upsample(flow3, 2)
        mask2 = Upsample(mask3, 2)
        warp2 = (flow2*self.scale/self.strides[4]).unsqueeze(1)
        warp2 = torch.repeat_interleave(warp2, 9, 1)
        S1, S2, S3, S4, S5 = warp2.shape
        warp2 = warp2.view(S1, S2*S3, S4, S5)
        warp2 = self.deform2(c22, warp2)
        tradeoff2 = feat2
        warp2 = (warp2 * F.sigmoid(mask2)) + self.conv2f(tradeoff2)
        warp2 = self.leakyRELU(warp2)
        corr2 = self.corr(c12, warp2)
        corr2 = self.leakyRELU(corr2)
        x = torch.cat((corr2, c12, feat2, flow2), 1)
        x = torch.cat((self.conv2_0(x), x),1)
        x = torch.cat((self.conv2_1(x), x),1)
        x = torch.cat((self.conv2_2(x), x),1)
        x = torch.cat((self.conv2_3(x), x),1)
        x = torch.cat((self.conv2_4(x), x),1)
        flow2 = flow2 + self.pred_flow2(x)

        x = self.dc_conv4(self.dc_conv3(self.dc_conv2(self.dc_conv1(x))))
        flow2 = flow2 + self.dc_conv7(self.dc_conv6(self.dc_conv5(x)))

        predictions = [flow * self.scale for flow in [flow6, flow5, flow4, flow3, flow2]]
        occlusion_masks = []
        occlusion_masks.append(F.sigmoid(mask2))
        c1s = [c11, c12, c13, c14, c15, c16]
        c2s = [c21, c12, c13, c24, c25, c26]
        flows = [flow6, flow5, flow4, flow3, flow2]
        mask0 = Upsample(mask2, 4)
        mask0 = F.sigmoid(mask0) - 0.5
        c30 = im1
        c40 = self.warp(im2, Upsample(flow2, 4)*self.scale)
        c30 = torch.cat((c30, torch.zeros_like(mask0)), 1)
        c40 = torch.cat((c40, mask0), 1)
        srcs = [c1s, c2s, flows, c30, c40]
        return predictions, occlusion_masks, srcs

class MaskFlownet(nn.Module):
    def __init__(self, config, **kwargs):
        super(MaskFlownet, self).__init__(**kwargs)
        self.strides = [64, 32, 16, 8, 4]
        self.md = 2
        self.scale = 20. * config.network.flow_multiplier.get(1.)
        self.deform_bias = config.network.deform_bias.get(True)
        self.upfeat_ch = config.network.upfeat_ch.get([16, 16, 16, 16])

        self.MaskFlownet_S = MaskFlownet_S(config)
        self.activate = nn.LeakyReLU(0.1)

        self.conv1x = conv(4,   16, stride=2)
        self.conv1y = conv(16,  16, stride=1)
        self.conv1z = conv(16,  16, stride=1)
        self.conv2x = conv(16,  32, stride=2)
        self.conv2y = conv(32,  32, stride=1)
        self.conv2z = conv(32,  32, stride=1)
        self.conv3x = conv(32,  64, stride=2)
        self.conv3y = conv(64,  64, stride=1)
        self.conv3z = conv(64,  64, stride=1)
        self.conv4x = conv(64,  96, stride=2)
        self.conv4y = conv(96,  96, stride=1)
        self.conv4z = conv(96,  96, stride=1)
        self.conv5x = conv(96,  128, stride=2)
        self.conv5y = conv(128, 128, stride=1)
        self.conv5z = conv(128, 128, stride=1)
        self.conv6x = conv(128, 196, stride=2)
        self.conv6y = conv(196, 196, stride=1)
        self.conv6z = conv(196, 196, stride=1)

        self.leakyRELU = nn.LeakyReLU(0.1)
        self.corr    = Correlation(pad_size=self.md, kernel_size=1, max_displacement=self.md, stride1=1, stride2=1, corr_multiply=1)

        nd = (2*self.md+1)**2
        dd = np.cumsum([128,128,96,64,32])

        od = nd+nd+2
        self.conv6_0 = conv(od,       128, kernel_size=3, stride=1)
        self.conv6_1 = conv(od+dd[0], 128, kernel_size=3, stride=1)
        self.conv6_2 = conv(od+dd[1], 96,  kernel_size=3, stride=1)
        self.conv6_3 = conv(od+dd[2], 64,  kernel_size=3, stride=1)
        self.conv6_4 = conv(od+dd[3], 32,  kernel_size=3, stride=1)
        self.pred_flow6 = predict_flow(od + dd[4])
        self.upfeat5 = deconv(od+dd[4], self.upfeat_ch[0], kernel_size=4, stride=2, padding=1)

        # od = nd+128+4
        od = nd+nd+128+16+2+2
        self.conv5_0 = conv(od,      128, kernel_size=3, stride=1)
        self.conv5_1 = conv(od+dd[0],128, kernel_size=3, stride=1)
        self.conv5_2 = conv(od+dd[1],96,  kernel_size=3, stride=1)
        self.conv5_3 = conv(od+dd[2],64,  kernel_size=3, stride=1)
        self.conv5_4 = conv(od+dd[3],32,  kernel_size=3, stride=1)
        self.pred_flow5 = predict_flow(od + dd[4])
        self.upfeat4 = deconv(od+dd[4], self.upfeat_ch[1], kernel_size=4, stride=2, padding=1)
        # self.deconv5 = deconv(2, 2, kernel_size=4, stride=2, padding=1)
        # self.upfeat5 = deconv(od+dd[4], 2, kernel_size=4, stride=2, padding=1)

        # od = nd+96+4
        # od = nd+96+18
        od = nd+nd+96+16+2+2
        self.conv4_0 = conv(od,      128, kernel_size=3, stride=1)
        self.conv4_1 = conv(od+dd[0],128, kernel_size=3, stride=1)
        self.conv4_2 = conv(od+dd[1],96,  kernel_size=3, stride=1)
        self.conv4_3 = conv(od+dd[2],64,  kernel_size=3, stride=1)
        self.conv4_4 = conv(od+dd[3],32,  kernel_size=3, stride=1)
        self.pred_flow4 = predict_flow(od + dd[4])
        self.upfeat3 = deconv(od+dd[4], self.upfeat_ch[2], kernel_size=4, stride=2, padding=1)
        # self.deconv4 = deconv(2, 2, kernel_size=4, stride=2, padding=1)
        # self.upfeat4 = deconv(od+dd[4], 2, kernel_size=4, stride=2, padding=1)

        # od = nd+64+4
        # od = nd+64+18
        od = nd+nd+64+16+2+2
        self.conv3_0 = conv(od,      128, kernel_size=3, stride=1)
        self.conv3_1 = conv(od+dd[0],128, kernel_size=3, stride=1)
        self.conv3_2 = conv(od+dd[1],96,  kernel_size=3, stride=1)
        self.conv3_3 = conv(od+dd[2],64,  kernel_size=3, stride=1)
        self.conv3_4 = conv(od+dd[3],32,  kernel_size=3, stride=1)
        self.pred_flow3 = predict_flow(od + dd[4])
        self.upfeat2 = deconv(od+dd[4], self.upfeat_ch[3], kernel_size=4, stride=2, padding=1)
        # self.deconv3 = deconv(2, 2, kernel_size=4, stride=2, padding=1)
        # self.upfeat3 = deconv(od+dd[4], 2, kernel_size=4, stride=2, padding=1)

        # od = nd+32+4
        # od = nd+32+18
        od = nd+nd+32+16+2+2
        self.conv2_0 = conv(od,      128, kernel_size=3, stride=1)
        self.conv2_1 = conv(od+dd[0],128, kernel_size=3, stride=1)
        self.conv2_2 = conv(od+dd[1],96,  kernel_size=3, stride=1)
        self.conv2_3 = conv(od+dd[2],64,  kernel_size=3, stride=1)
        self.conv2_4 = conv(od+dd[3],32,  kernel_size=3, stride=1)
        self.pred_flow2 = predict_flow(od + dd[4])

        self.dc_conv1 = conv(od+dd[4], 128, kernel_size=3, stride=1, padding=1,  dilation=1)
        self.dc_conv2 = conv(128,      128, kernel_size=3, stride=1, padding=2,  dilation=2)
        self.dc_conv3 = conv(128,      128, kernel_size=3, stride=1, padding=4,  dilation=4)
        self.dc_conv4 = conv(128,      96,  kernel_size=3, stride=1, padding=8,  dilation=8)
        self.dc_conv5 = conv(96,       64,  kernel_size=3, stride=1, padding=16, dilation=16)
        self.dc_conv6 = conv(64,       32,  kernel_size=3, stride=1, padding=1,  dilation=1)
        self.dc_conv7 = predict_flow(32)

        # self.upfeat5 = deconv()

        self.deform6 = deformable_conv(196, 196)
        self.deform5 = deformable_conv(128, 128)
        self.deform4 = deformable_conv(96, 96)
        self.deform3 = deformable_conv(64, 64)
        self.deform2 = deformable_conv(32, 32)

    def warp(self, x, flo):
        """
        warp an image/tensor (im2) back to im1, according to the optical flow
        x: [B, C, H, W] (im2)
        flo: [B, 2, H, W] flow
        """
        B, C, H, W = x.size()
        # mesh grid
        xx = torch.arange(0, W).view(1,-1).repeat(H,1)
        yy = torch.arange(0, H).view(-1,1).repeat(1,W)
        xx = xx.view(1,1,H,W).repeat(B,1,1,1)
        yy = yy.view(1,1,H,W).repeat(B,1,1,1)
        grid = torch.cat((xx,yy),1).float()

        if x.is_cuda:
            grid = grid.cuda()
        # vgrid = Variable(grid) + flo
        vgrid = Variable(grid) + torch.flip(flo, [1])

        # scale grid to [-1,1]
        vgrid[:,0,:,:] = 2.0*vgrid[:,0,:,:].clone() / max(W-1,1)-1.0
        vgrid[:,1,:,:] = 2.0*vgrid[:,1,:,:].clone() / max(H-1,1)-1.0

        # vgrid = vgrid.permute(0,2,3,1)
        vgrid = vgrid.permute(0,2,3,1).clamp(-1.1, 1.1)
        output = nn.functional.grid_sample(x, vgrid, align_corners=True)
        mask = torch.autograd.Variable(torch.ones(x.size())).cuda()
        mask = nn.functional.grid_sample(mask, vgrid, align_corners=True)

        # if W==128:
        # np.save('mask.npy', mask.cpu().data.numpy())
        # np.save('warp.npy', output.cpu().data.numpy())

        mask[mask<0.9999] = 0
        mask[mask>0] = 1

        return output*mask

    def forward(self, im1, im2):

        # im1 = x[:,:3,:,:]
        # im2 = x[:,3:,:,:]

        _, _, srcs = self.MaskFlownet_S(im1, im2)
        c1s, c2s, flows, c30, c40 = srcs
        c11, c12, c13, c14, c15, c16 = c1s
        c21, c22, c23, c24, c25, c26 = c2s

        c31 = self.conv1z(self.conv1y(self.conv1x(c30)))
        c32 = self.conv2z(self.conv2y(self.conv2x(c31)))
        c33 = self.conv3z(self.conv3y(self.conv3x(c32)))
        c34 = self.conv4z(self.conv4y(self.conv4x(c33)))
        c35 = self.conv5z(self.conv5y(self.conv5x(c34)))
        c36 = self.conv6z(self.conv6y(self.conv6x(c35)))

        c41 = self.conv1z(self.conv1y(self.conv1x(c40)))
        c42 = self.conv2z(self.conv2y(self.conv2x(c41)))
        c43 = self.conv3z(self.conv3y(self.conv3x(c42)))
        c44 = self.conv4z(self.conv4y(self.conv4x(c43)))
        c45 = self.conv5z(self.conv5y(self.conv5x(c44)))
        c46 = self.conv6z(self.conv6y(self.conv6x(c45)))

        flow6 = flows[0]

        warp6u = (flow6*self.scale/self.strides[0]).unsqueeze(1)
        warp6u = torch.repeat_interleave(warp6u, 9, 1)
        S1, S2, S3, S4, S5 = warp6u.shape
        warp6u = warp6u.view(S1, S2*S3, S4, S5)
        warp6u = self.deform6(c26, warp6u)
        warp6u = self.leakyRELU(warp6u)
        corr6u = self.leakyRELU(self.corr(c16, warp6u))
        warp6v = c46
        corr6v = self.leakyRELU(self.corr(c36, warp6v))
        x = torch.cat((corr6u, corr6v, flow6),1)
        x = torch.cat((self.conv6_0(x), x),1)
        x = torch.cat((self.conv6_1(x), x),1)
        x = torch.cat((self.conv6_2(x), x),1)
        x = torch.cat((self.conv6_3(x), x),1)
        x = torch.cat((self.conv6_4(x), x),1)
        flow6 = flow6 + self.pred_flow6(x)

        feat5 = self.leakyRELU(self.upfeat5(x))
        flow5 = Upsample(flow6, 2)
        warp5u = (flow5*self.scale/self.strides[1]).unsqueeze(1)
        warp5u = torch.repeat_interleave(warp5u, 9, 1)
        S1, S2, S3, S4, S5 = warp5u.shape
        warp5u = warp5u.view(S1, S2*S3, S4, S5)
        warp5u = self.deform5(c25, warp5u)
        warp5u = self.leakyRELU(warp5u)
        corr5u = self.leakyRELU(self.corr(c15, warp5u))
        warp5v = c45
        corr5v = self.leakyRELU(self.corr(c35, warp5v))
        x = torch.cat((c15, feat5, corr5u, corr5v, flow5, flows[1]),1)
        x = torch.cat((self.conv5_0(x), x),1)
        x = torch.cat((self.conv5_1(x), x),1)
        x = torch.cat((self.conv5_2(x), x),1)
        x = torch.cat((self.conv5_3(x), x),1)
        x = torch.cat((self.conv5_4(x), x),1)
        flow5 = flow5 + self.pred_flow5(x)

        feat4 = self.leakyRELU(self.upfeat4(x))
        flow4 = Upsample(flow5, 2)
        warp4u = (flow4*self.scale/self.strides[2]).unsqueeze(1)
        warp4u = torch.repeat_interleave(warp4u, 9, 1)
        S1, S2, S3, S4, S5 = warp4u.shape
        warp4u = warp4u.view(S1, S2*S3, S4, S5)
        warp4u = self.deform4(c24, warp4u)
        warp4u = self.leakyRELU(warp4u)
        corr4u = self.leakyRELU(self.corr(c14, warp4u))
        warp4v = c44
        corr4v = self.leakyRELU(self.corr(c34, warp4v))
        x = torch.cat((c14, feat4, corr4u, corr4v, flow4, flows[2]),1)
        x = torch.cat((self.conv4_0(x), x),1)
        x = torch.cat((self.conv4_1(x), x),1)
        x = torch.cat((self.conv4_2(x), x),1)
        x = torch.cat((self.conv4_3(x), x),1)
        x = torch.cat((self.conv4_4(x), x),1)
        flow4 = flow4 + self.pred_flow4(x)

        feat3 = self.leakyRELU(self.upfeat3(x))
        flow3 = Upsample(flow4, 2)
        warp3u = (flow3*self.scale/self.strides[3]).unsqueeze(1)
        warp3u = torch.repeat_interleave(warp3u, 9, 1)
        S1, S2, S3, S4, S5 = warp3u.shape
        warp3u = warp3u.view(S1, S2*S3, S4, S5)
        warp3u = self.deform3(c23, warp3u)
        warp3u = self.leakyRELU(warp3u)
        corr3u = self.leakyRELU(self.corr(c13, warp3u))
        warp3v = c43
        corr3v = self.leakyRELU(self.corr(c33, warp3v))
        x = torch.cat((c13, feat3, corr3u, corr3v, flow3, flows[3]),1)
        x = torch.cat((self.conv3_0(x), x),1)
        x = torch.cat((self.conv3_1(x), x),1)
        x = torch.cat((self.conv3_2(x), x),1)
        x = torch.cat((self.conv3_3(x), x),1)
        x = torch.cat((self.conv3_4(x), x),1)
        flow3 = flow3 + self.pred_flow3(x)

        feat2 = self.leakyRELU(self.upfeat2(x))
        flow2 = Upsample(flow3, 2)
        warp2u = (flow2*self.scale/self.strides[4]).unsqueeze(1)
        warp2u = torch.repeat_interleave(warp2u, 9, 1)
        S1, S2, S3, S4, S5 = warp2u.shape
        warp2u = warp2u.view(S1, S2*S3, S4, S5)
        warp2u = self.deform2(c22, warp2u)
        warp2u = self.leakyRELU(warp2u)
        corr2u = self.leakyRELU(self.corr(c12, warp2u))
        warp2v = c42
        corr2v = self.leakyRELU(self.corr(c32, warp2v))
        x = torch.cat((c12, feat2, corr2u, corr2v, flow2, flows[4]),1)
        x = torch.cat((self.conv2_0(x), x),1)
        x = torch.cat((self.conv2_1(x), x),1)
        x = torch.cat((self.conv2_2(x), x),1)
        x = torch.cat((self.conv2_3(x), x),1)
        x = torch.cat((self.conv2_4(x), x),1)
        flow2 = flow2 + self.pred_flow2(x)

        x = self.dc_conv4(self.dc_conv3(self.dc_conv2(self.dc_conv1(x))))
        flow2 = flow2 + self.dc_conv7(self.dc_conv6(self.dc_conv5(x)))

        preds = [flow * self.scale for flow in [flow6, flow5, flow4, flow3, flow2]]
        visuals = []
        visuals.append(flow2[:,:1])
        return preds, visuals, []

class EpeLoss(nn.Module):
    def __init__(self, eps = 0):
        super(EpeLoss, self).__init__()
        self.eps = eps

    def forward(self, pred, label):
        loss = ((pred - label).pow(2).sum(1) + self.eps).sqrt()
        return loss.view(loss.shape[0], -1).mean(1)

class EpeLossWithMask(nn.Module):
    def __init__(self, eps=1e-8, q=None):
        super(EpeLossWithMask, self).__init__()
        self.eps = eps
        self.q = q

    def forward(self, pred, label, mask):
        if self.q is not None:
            loss = ((pred - label).abs().sum(1) + self.eps) ** self.q
        else:
            loss = ((pred - label).pow(2).sum(1) + self.eps).sqrt()
        loss = loss * mask.squeeze(1)
        loss = loss.view(loss.shape[0], -1).sum(1) / mask.view(mask.shape[0], -1).sum(1)
        return loss

class MultiscaleEpe(nn.Module):
    def __init__(self, scales, weights, match, eps = 1e-8, q = None):
        super(MultiscaleEpe, self).__init__()

        self.scales = scales
        self.weights = weights
        self.match = match
        self.eps = eps
        self.q = q

    def forward(self, flow, mask, *predictions):
        losses = 0
        if self.match == 'upsampling':
            for p, w, s in zip(predictions, self.weights, self.scales):
                losses += EpeLossWithMask(eps=self.eps, q=self.q)(Upsample(p, s), flow, mask) * w
        elif self.match == 'downsampling':
            for p, w, s in zip(predictions, self.weights, self.scales):
                losses += EpeLossWithMask(eps=self.eps, q=self.q)(p, Downsample(flow, s), Downsample(mask, s)) * w
        else:
            raise NotImplementedError
        return losses