import os
import tensorflow as tf
import numpy as np
from math import exp
from tqdm import tqdm
from nltk.translate.bleu_score import corpus_bleu, SmoothingFunction
from pycocoevalcap.meteor.meteor import Meteor
from pycocoevalcap.cider.cider import Cider
from pycocoevalcap.rouge.rouge import Rouge
from summary_handler import SummaryHandler
from read_data import *
from data import GenModelVocab, translate, save_vocab, restore_vocab,\
    translate_spans, GloVEVocab
import time

from tensorflow.python.client import timeline

from elmo import Batcher, BidirectionalLanguageModel

from utils import *

import sys

import random

from commonsense_utils.commonsense_data import COMMONSENSE_REL_LOOKUP

VALID_VERSIONS = ['baseline_nqa', 'commonsense_nqa', 'baseline_wh',
    'commonsense_wh']

VALID_MODES = ['train', 'test', 'get-vocab', 'build_dataset', 'build_wikihop_dataset']

def import_model(config):
    global UsedModel
    if config.version == 'baseline_nqa':
        from model_baseline_nqa import\
            NarrativeQAGatedMultiBidafWithSelfAttnAndELMoCUDNN as\
            UsedModel
    elif config.version == 'commonsense_nqa':
        from model_commonsense_nqa import \
            GatedMultiBidaf_ResSA_ELMo_PtrGen_CS_1_fix_path_concat_project_w_rels\
            as UsedModel
    elif config.version == 'baseline_wh':
        from model_baseline_wh import \
            GatedMultiBidafWithSelfAttnAndELMoCUDNN_MC as UsedModel
    elif config.version == 'commonsense_wh':
        from model_commonsense_wh import \
            GatedMultiBidaf_ResSA_ELMo_PtrGen_CS_path_cp_rels_MC as\
            UsedModel
    else:
        raise ValueError("Model version DNE!")

def main(config):
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # disable cpp error msgs
    if config.mode == 'train':
        _train(config)
    elif config.mode == 'test':
        _test(config)
    elif config.mode == 'get-vocab':
        _print_vocab(config)
    elif config.mode == 'build_dataset':
        _build_datasets(config)
    elif config.mode == 'build_wikihop_dataset':
        _build_wikihop_datasets(config)
    elif config.mode == 'generate_answers':
        _generate_answers(config)
    else:
        raise ValueError("Invalid Mode!")

def _print_vocab(config):
    train_data = load_processed_dataset(config, 'train')
    vocabs = train_data.get_word_lists()
    valid_data = load_processed_dataset(config, 'valid')
    vocabs.update(valid_data.get_word_lists())
    for v in vocabs:
        print(v.encode('utf-8'))

def _build_datasets(config):
    train_data_processed = create_processed_dataset(config, 'train')
    valid_data_processed = create_processed_dataset(config, 'valid')
    test_data_processed = create_processed_dataset(config, 'test')

    save_processed_dataset(config, train_data_processed, 'train')
    save_processed_dataset(config, valid_data_processed, 'valid')
    save_processed_dataset(config, test_data_processed, 'test')

def _build_wikihop_datasets(config):
    train_data_processed, dev_data_processed = create_processed_wikihop_dataset_cs(config)
    save_wikihop_processed_dataset(config, train_data_processed, 'train')
    save_wikihop_processed_dataset(config, dev_data_processed, 'valid')

def _train(config):
    import_model(config)

    train_data = load_processed_dataset(config, 'train')
    train_commonsense_list = []

    vocab_freq = train_data.get_word_lists()
    # Fixme: values or keys?
    # relations_vocab = COMMONSENSE_REL_LOOKUP.values()
    relations_vocab = COMMONSENSE_REL_LOOKUP.keys()

    valid_data = load_processed_dataset(config, 'valid')

    print("Data loaded!")
    vocab = GenModelVocab(vocab_freq, config.embedding_size,
            forced_word_list=relations_vocab,
            cs=train_commonsense_list,
            threshold=config.min_occurence)

    print("Vocab built! Size (%d)" % vocab.size())

    model = UsedModel(config, vocab)

    #create session
    gpu_configuration = gpu_config()
    sess = tf.Session(config=gpu_configuration)
    with sess.as_default():
        model.build_graph()
        print("Graph built!")
        model.add_train_op()
        print("Train op added!")
    sess.run(tf.global_variables_initializer())
    print("Variables initialized")

    if config.continue_training:
        start_e, steps, out_dir, ckpt_dir = restore_from_last_ckpt(
                config, model, sess)

        # backup new argv
        with open(os.path.join(out_dir, 'argv.txt'), 'a') as f:
            f.write("\n")
            f.write(" ".join(sys.argv))

        print("Continue training after epoch %d, step %d" % (start_e, steps))

    else:
        if config.model_name == 'default':
            c_time = time.strftime("%m_%d_%H_%M_%S", time.localtime())
            commonsense = ""
            if config.load_commonsense:
                commonsense = "_with_cs"

            if config.shuffle_cs:
                commonsense = "_shuffle_CS"

            config.model_name = UsedModel.__name__ + commonsense + "_%s" % c_time

        if config.debug:
            config.checkpoint_size = 10

        if not config.debug:
            # XXX factor into class eventually
            out_dir = os.path.join(config.out_root, config.model_name)
            if os.path.exists(out_dir):
                raise ValueError("Output directory already exists!")
            else:
                os.makedirs(out_dir)

            # back up src file
            os.system("cp -r src %s" % os.path.join(out_dir, 'src'))
            # back up argv
            with open(os.path.join(out_dir, "argv.txt"), 'w') as f:
                f.write(" ".join(sys.argv))

            # back up environ
            with open(os.path.join(out_dir, 'recreate_environ.sh'), 'w') as f:
                for var, val in os.environ.items():
                    f.write("export %s=\"%s\"\n" % (var, val))

            os.system("chmod +x %s" % os.path.join(out_dir, 'recreate_environ.sh'))

            ckpt_dir = os.path.join(out_dir, "ckpts")

            vocab_loc = os.path.join(out_dir, "vocab.pkl")
            save_vocab(vocab, vocab_loc)

            print("Initialized output at %s" % out_dir)

        steps = 0
        start_e = -1

        print("Started training!")

    #construct graph handler
    if not config.multiple_choice:
        summary_handler = SummaryHandler(
                os.path.join(config.summary_save_path, config.model_name),
                ['LOSS', 'PERPLEXITY', 'EVAL_PERPLEXITY', 'BLEU_1', 'BLEU_4',
                    'METEOR', 'ROUGE', 'CIDER'])
    else:
        summary_handler = SummaryHandler(
                os.path.join(config.summary_save_path, config.model_name),
                ['LOSS', 'ACCURACY'])

    for e in range(config.num_epochs):
        total_loss = []
        for batches in tqdm(train_data.get_batches(config.batch_size,
            front_heavy = (e == 0))):

            if steps != 0 or not config.start_eval:
                steps += 1

            if steps > 10 and config.debug:
                exit(0)

            is_training = True

            fd = model.encode(batches, is_training)
            oovs = model.get_batch_oov()

            loss = model.train_step(sess, fd)
            total_loss.append(loss)

            if steps % config.checkpoint_size == 0:
                if not config.multiple_choice:
                    bleu1, bleu4, meteor, rouge, cider, eval_loss, preds = eval_dataset(
                            config, valid_data, vocab, model, sess)

                    print("Result at step %d:" % steps)
                    print("Bleu1: ", bleu1)
                    print("Bleu4: ", bleu4)
                    print("Meteor: ", meteor)
                    print("Rouge-L: ", rouge)
                    print("CIDEr: ", cider)

                    if not config.debug:
                        write_summaries(sess, summary_handler, avg(total_loss),
                                eval_loss, bleu1, bleu4, meteor, rouge, cider, steps)

                        if start_e > 0:
                            epoch = e + start_e
                        else:
                            epoch = e

                        model.save_to(sess, os.path.join(ckpt_dir,
                            "epoch_%04d_step%08d_bleu1(%f)_bleu4(%f)_meteor(%f)_rogue(%f)" %
                            (epoch, steps, bleu1, bleu4, meteor, rouge)))
                        print("Model saved!")
                else:
                    accuracy = eval_multiple_choice_dataset(
                            config, valid_data, vocab, model, sess)
                    print("Result at step %d: %f" % (steps, accuracy))

                    if not config.debug:
                        summary_handler.write_summaries(sess,
                                {
                                    'ITERATION': steps,
                                    'LOSS': avg(total_loss),
                                    'ACCURACY': accuracy
                                })

                        if start_e > 0:
                            epoch = e + start_e
                        else:
                            epoch = e

                        model.save_to(sess, os.path.join(ckpt_dir,
                            'epoch_%04d_step%08d_acc(%f)' % (epoch, steps,
                                accuracy)))

    summary_handler.close_writer()

def _test(config):
    print("Evaluating!")
    import_model(config)

    out_dir = os.path.join(config.out_root, config.model_name)

    vocab_loc = os.path.join(out_dir, "vocab.pkl")
    if os.path.exists(vocab_loc): # vocab exists!
        vocab = restore_vocab(vocab_loc)
    else:
        raise Exception("Not valid output directory! No vocab found!")

    print("Vocab built! Size (%d)" % vocab.size())

    if config.use_dev:
        valid_data = load_processed_dataset(config, 'valid')
    else:
        print("Using test!")
        valid_data = load_processed_dataset(config, 'test')

    if config.shuffle_cs:
        valid_data = shuffle_cs(valid_data)

    print("Data loaded!")

    #construct model
    model = UsedModel(config, vocab)

    gpu_configuration = gpu_config()
    sess = tf.Session(config=gpu_configuration)
    with sess.as_default():
        model.build_graph()
        print("Graph built!")
        model.add_train_op()
        print("Train op added!")

    sess.run(tf.global_variables_initializer())

    if config.use_ckpt is not None:
        model.restore_from(sess, os.path.join(out_dir, 'ckpts', config.use_ckpt))
    elif config.at_step is not None:
        restore_from_step(config, model, sess, config.at_step)
    else:
        raise ValueError("Must specify a ckpt to restore from!")

    if not config.multiple_choice:
        bleu1, bleu4, meteor, rouge, cider, eval_loss, preds = eval_dataset(
                config, valid_data, vocab, model, sess)

        #print("Result for %s" % config.use_ckpt)
        #print("Bleu1: ", bleu1)
        #print("Bleu4: ", bleu4)
        #print("Meteor: ", meteor)
        #print("Rouge-L: ", rouge)
        #print("CIDEr: ", cider)

        filename = config.model_name + '_preds.txt'
        with open(filename, 'w') as pred_file:
            for p in preds:
                pred_file.write(" ".join(p) + "\n")
    else:
        accuracy = eval_multiple_choice_dataset( config, valid_data, vocab,
                model, sess)
        print("Result for %s: %f" % (config.use_ckpt, accuracy))

    print("Done!")

def _generate_answers(config):

    def save_answers(prefix, dset):
        ref0 = open(prefix + '_ref0.txt', 'w')
        ref1 = open(prefix + '_ref1.txt', 'w')

        punc_list = ['.', '!', '?']

        for item in dset.data:
            ref0.write(" ".join(filter(lambda w: w not in punc_list,
                item['answer1'])) + '\n')
            ref1.write(" ".join(filter(lambda w: w not in punc_list,
                item['answer2'])) + '\n')

        ref0.close()
        ref1.close()

    val_data = load_processed_dataset(config, 'valid')
    test_data = load_processed_dataset(config, 'test')

    save_answers('val', val_data)
    save_answers('test', test_data)


"""Copy number detection with CNVkit with specific support for targeted sequencing.

http://cnvkit.readthedocs.org
"""
import collections
import copy
import glob
import math
import operator
import os
import shutil
import sys
import tempfile

import pybedtools
import numpy as np
import toolz as tz

from bcbio import utils
from bcbio.bam import ref
from bcbio.distributed.multi import run_multicore, zeromq_aware_logging
from bcbio.distributed.transaction import file_transaction
from bcbio.heterogeneity import chromhacks
from bcbio.pipeline import datadict as dd
from bcbio.pipeline import config_utils
from bcbio.provenance import do
from bcbio.variation import effects, ploidy, population, vcfutils
from bcbio.structural import annotate, plot, shared
from functools import reduce

def use_general_sv_bins(data):
    """Check if we should use a general binning approach for a sample.

    Checks if CNVkit is enabled and we haven't already run CNVkit.
    """
    if any([c in dd.get_svcaller(data) for c in ["cnvkit", "titancna", "purecn", "gatk-cnv"]]):
        if not _get_original_coverage(data):
            return True
    return False

def bin_approach(data):
    """Check for binning approach from configuration or normalized file.
    """
    for approach in ["cnvkit", "gatk-cnv"]:
        if approach in dd.get_svcaller(data):
            return approach
    norm_file = tz.get_in(["depth", "bins", "normalized"], data)
    if norm_file.endswith(("-crstandardized.tsv", "-crdenoised.tsv")):
        return "gatk-cnv"
    if norm_file.endswith(".cnr"):
        return "cnvkit"

def run(items, background=None):
    """Detect copy number variations from batched set of samples using CNVkit.
    """
    if not background: background = []
    return _cnvkit_by_type(items, background)

def _sv_workdir(data):
    return utils.safe_makedir(os.path.join(data["dirs"]["work"], "structural",
                                           dd.get_sample_name(data), "cnvkit"))

def _cnvkit_by_type(items, background):
    """Dispatch to specific CNVkit functionality based on input type.
    """
    if len(items + background) == 1:
        return _run_cnvkit_single(items[0])
    elif vcfutils.get_paired_phenotype(items[0]):
        return _run_cnvkit_cancer(items, background)
    else:
        return _run_cnvkit_population(items, background)

def _associate_cnvkit_out(ckouts, items, is_somatic=False):
    """Associate cnvkit output with individual items.
    """
    assert len(ckouts) == len(items)
    out = []
    upload_counts = collections.defaultdict(int)
    for ckout, data in zip(ckouts, items):
        ckout = copy.deepcopy(ckout)
        ckout["variantcaller"] = "cnvkit"
        if utils.file_exists(ckout["cns"]) and _cna_has_values(ckout["cns"]):
            ckout = _add_seg_to_output(ckout, data)
            ckout = _add_gainloss_to_output(ckout, data)
            ckout = _add_segmetrics_to_output(ckout, data)
            ckout = _add_variantcalls_to_output(ckout, data, items, is_somatic)
            # ckout = _add_coverage_bedgraph_to_output(ckout, data)
            ckout = _add_cnr_bedgraph_and_bed_to_output(ckout, data)
            if "svplots" in dd.get_tools_on(data):
                ckout = _add_plots_to_output(ckout, data)
            ckout["do_upload"] = upload_counts[ckout.get("vrn_file")] == 0
        if "sv" not in data:
            data["sv"] = []
        data["sv"].append(ckout)
        if ckout.get("vrn_file"):
            upload_counts[ckout["vrn_file"]] += 1
        out.append(data)
    return out

def _run_cnvkit_single(data, background=None):
    """Process a single input file with BAM or uniform background.
    """
    if not background:
        background = []
    ckouts = _run_cnvkit_shared([data], background)
    if not ckouts:
        return [data]
    else:
        assert len(ckouts) == 1
        return _associate_cnvkit_out(ckouts, [data])

def _run_cnvkit_cancer(items, background):
    """Run CNVkit on a tumor/normal pair.
    """
    paired = vcfutils.get_paired_bams([x["align_bam"] for x in items], items)
    normal_data = [x for x in items if dd.get_sample_name(x) != paired.tumor_name]
    tumor_ready, normal_ready = _match_batches(paired.tumor_data, normal_data[0] if normal_data else None)
    ckouts = _run_cnvkit_shared([tumor_ready], [normal_ready] if normal_ready else [])
    if not ckouts:
        return items
    assert len(ckouts) == 1
    tumor_data = _associate_cnvkit_out(ckouts, [paired.tumor_data], is_somatic=True)
    return tumor_data + normal_data

def _match_batches(tumor, normal):
    """Fix batch names for shared tumor/normals to ensure matching
     """
    def _get_batch(x):
        b = dd.get_batch(x)
        return [b] if not isinstance(b, (list, tuple)) else b
    if normal:
        tumor = copy.deepcopy(tumor)
        normal = copy.deepcopy(normal)
        cur_batch = list(set(_get_batch(tumor)) & set(_get_batch(normal)))
        assert len(cur_batch) == 1, "No batch overlap: %s and %s" % (_get_batch(tumor), _get_batch(normal))
        cur_batch = cur_batch[0]
        tumor["metadata"]["batch"] = cur_batch
        normal["metadata"]["batch"] = cur_batch
    return tumor, normal

def _run_cnvkit_population(items, background):
    """Run CNVkit on a population of samples.

    Tries to calculate background based on case/controls, otherwise
    uses samples from the same batch as background.
    """
    if background and len(background) > 0:
        inputs = items
    else:
        inputs, background = shared.find_case_control(items)

    # if we have case/control organized background or a single sample
    if len(inputs) == 1 or len(background) > 0:
        ckouts = _run_cnvkit_shared(inputs, background)
        return _associate_cnvkit_out(ckouts, inputs) + background
    # otherwise run each sample with the others in the batch as background
    else:
        out = []
        for cur_input in items:
            background = [d for d in items if dd.get_sample_name(d) != dd.get_sample_name(cur_input)]
            ckouts = _run_cnvkit_shared([cur_input], background)
            out.extend(_associate_cnvkit_out(ckouts, [cur_input]))
        return out

def _get_cmd(script_name="cnvkit.py"):
    return os.path.join(os.path.dirname(os.path.realpath(sys.executable)), script_name)

def _prep_cmd(cmd, tx_out_file):
    """Wrap CNVkit commands ensuring we use local temporary directories.
    """
    cmd = " ".join(cmd) if isinstance(cmd, (list, tuple)) else cmd
    return "export TMPDIR=%s && %s" % (os.path.dirname(tx_out_file), cmd)

def _bam_to_outbase(bam_file, work_dir, data):
    """Convert an input BAM file into CNVkit expected output.

    Handles previous non-batch cases to avoid re-calculating,
    returning both new and old values:
    """
    batch = dd.get_batch(data) or dd.get_sample_name(data)
    out_base = os.path.splitext(os.path.basename(bam_file))[0].split(".")[0]
    base = os.path.join(work_dir, out_base)
    return "%s-%s" % (base, batch), base

def _get_original_coverage(data, itype="target"):
    """Back compatible: get existing coverage files if they exist
    """
    work_dir = os.path.join(_sv_workdir(data), "raw")
    work_bam = dd.get_work_bam(data) or dd.get_align_bam(data)
    out = []
    base, _ = _bam_to_outbase(work_bam, work_dir, data)
    target_cnn = "%s.targetcoverage.cnn" % base
    anti_cnn = "%s.antitargetcoverage.cnn" % base
    if os.path.exists(target_cnn) and os.path.exists(anti_cnn):
        out.append({"bam": work_bam, "file": target_cnn, "cnntype": "target",
                    "itype": itype, "sample": dd.get_sample_name(data)})
        out.append({"bam": work_bam, "file": anti_cnn, "cnntype": "antitarget",
                    "itype": itype, "sample": dd.get_sample_name(data)})
    return out

def _run_cnvkit_shared(inputs, backgrounds):
    """Shared functionality to run CNVkit, parallelizing over multiple BAM files.

    Handles new style cases where we have pre-normalized inputs and
    old cases where we run CNVkit individually.
    """
    if tz.get_in(["depth", "bins", "normalized"], inputs[0]):
        ckouts = []
        for data in inputs:
            cnr_file = tz.get_in(["depth", "bins", "normalized"], data)
            cns_file = os.path.join(_sv_workdir(data), "%s.cns" % dd.get_sample_name(data))
            cns_file = _cnvkit_segment(cnr_file, dd.get_coverage_interval(data), data,
                                       inputs + backgrounds, cns_file)
            ckouts.append({"cnr": cnr_file, "cns": cns_file,
                           "background": tz.get_in(["depth", "bins", "background"], data)})
        return ckouts
    else:
        return _run_cnvkit_shared_orig(inputs, backgrounds)

def _get_original_targets(data):
    """Back compatible: get pre-existing target BEDs.
    """
    work_dir = os.path.join(_sv_workdir(data), "raw")
    batch = dd.get_batch(data) or dd.get_sample_name(data)
    return (glob.glob(os.path.join(work_dir, "*-%s.target.bed" % batch))[0],
            glob.glob(os.path.join(work_dir, "*-%s.antitarget.bed" % batch))[0])

def _get_general_coverage(data, itype):
    """Retrieve coverage information from new shared SV bins.
    """
    work_bam = dd.get_align_bam(data) or dd.get_work_bam(data)
    return [{"bam": work_bam, "file": tz.get_in(["depth", "bins", "target"], data),
             "cnntype": "target", "itype": itype, "sample": dd.get_sample_name(data)},
            {"bam": work_bam, "file": tz.get_in(["depth", "bins", "antitarget"], data),
             "cnntype": "antitarget", "itype": itype, "sample": dd.get_sample_name(data)}]

def _run_cnvkit_shared_orig(inputs, backgrounds):
    """Original CNVkit implementation with full normalization and segmentation.
    """
    work_dir = _sv_workdir(inputs[0])
    raw_work_dir = utils.safe_makedir(os.path.join(work_dir, "raw"))
    background_name = dd.get_sample_name(backgrounds[0]) if backgrounds else "flat"
    background_cnn = os.path.join(raw_work_dir, "%s_background.cnn" % (background_name))
    ckouts = []
    for cur_input in inputs:
        cur_raw_work_dir = utils.safe_makedir(os.path.join(_sv_workdir(cur_input), "raw"))
        out_base, out_base_old = _bam_to_outbase(dd.get_align_bam(cur_input), cur_raw_work_dir, cur_input)
        if utils.file_exists(out_base_old + ".cns"):
            out_base = out_base_old
        ckouts.append({"cnr": "%s.cnr" % out_base,
                       "cns": "%s.cns" % out_base})
    if not utils.file_exists(ckouts[0]["cns"]):
        cov_interval = dd.get_coverage_interval(inputs[0])
        samples_to_run = list(zip(["background"] * len(backgrounds), backgrounds)) + \
                         list(zip(["evaluate"] * len(inputs), inputs))
        # New style shared SV bins
        if tz.get_in(["depth", "bins", "target"], inputs[0]):
            target_bed = tz.get_in(["depth", "bins", "target"], inputs[0])
            antitarget_bed = tz.get_in(["depth", "bins", "antitarget"], inputs[0])
            raw_coverage_cnns = reduce(operator.add,
                                       [_get_general_coverage(cdata, itype) for itype, cdata in samples_to_run])
        # Back compatible with pre-existing runs
        else:
            target_bed, antitarget_bed = _get_original_targets(inputs[0])
            raw_coverage_cnns = reduce(operator.add,
                                       [_get_original_coverage(cdata, itype) for itype, cdata in samples_to_run])
        # Currently metrics not calculated due to speed and needing re-evaluation
        # We could re-enable with larger truth sets to evaluate background noise
        # But want to reimplement in a more general fashion as part of normalization
        if False:
            coverage_cnns = reduce(operator.add,
                                [_cnvkit_metrics(cnns, target_bed, antitarget_bed, cov_interval,
                                                    inputs + backgrounds)
                                    for cnns in tz.groupby("bam", raw_coverage_cnns).values()])
            background_cnn = cnvkit_background(_select_background_cnns(coverage_cnns),
                                                background_cnn, inputs, target_bed, antitarget_bed)
        else:
            coverage_cnns = raw_coverage_cnns
            background_cnn = cnvkit_background([x["file"] for x in coverage_cnns if x["itype"] == "background"],
                                                background_cnn, inputs, target_bed, antitarget_bed)
        parallel = {"type": "local", "cores": dd.get_cores(inputs[0]), "progs": ["cnvkit"]}
        fixed_cnrs = run_multicore(_cnvkit_fix,
                                   [(cnns, background_cnn, inputs, ckouts) for cnns in
                                    tz.groupby("bam", [x for x in coverage_cnns
                                                       if x["itype"] == "evaluate"]).values()],
                                   inputs[0]["config"], parallel)
        [_cnvkit_segment(cnr, cov_interval, data, inputs + backgrounds) for cnr, data in fixed_cnrs]
    return ckouts

def _cna_has_values(fname):
    with open(fname) as in_handle:
        for i, line in enumerate(in_handle):
            if i > 0:
                return True
    return False

def _cnvkit_segment(cnr_file, cov_interval, data, items, out_file=None, detailed=False):
    """Perform segmentation and copy number calling on normalized inputs
    """
    if not out_file:
        out_file = "%s.cns" % os.path.splitext(cnr_file)[0]
    if not utils.file_uptodate(out_file, cnr_file):
        with file_transaction(data, out_file) as tx_out_file:
            if not _cna_has_values(cnr_file):
                with open(tx_out_file, "w") as out_handle:
                    out_handle.write("chromosome\tstart\tend\tgene\tlog2\tprobes\tCN1\tCN2\tbaf\tweight\n")
            else:
                # Scale cores to avoid memory issues with segmentation
                # https://github.com/etal/cnvkit/issues/346
                if cov_interval == "genome":
                    cores = max(1, dd.get_cores(data) // 2)
                else:
                    cores = dd.get_cores(data)
                cmd = [_get_cmd(), "segment", "-p", str(cores), "-o", tx_out_file, cnr_file]
                small_vrn_files = _compatible_small_variants(data, items)
                if len(small_vrn_files) > 0 and _cna_has_values(cnr_file) and cov_interval != "genome":
                    cmd += ["--vcf", small_vrn_files[0].name, "--sample-id", small_vrn_files[0].sample]
                    if small_vrn_files[0].normal:
                        cmd += ["--normal-id", small_vrn_files[0].normal]
                resources = config_utils.get_resources("options", data["config"])
                user_options = resources.get("cnvkit_segment", [])
                cmd += [str(x) for x in user_options]
                if cov_interval == "genome" and "--threshold" not in user_options:
                    cmd += ["--threshold", "0.00001"]
                # For tumors, remove very low normalized regions, avoiding upcaptured noise
                # https://github.com/bcbio/bcbio-nextgen/issues/2171#issuecomment-348333650
                # unless we want detailed segmentation for downstream tools
                paired = vcfutils.get_paired(items)
                if paired:
                    #if detailed:
                    #    cmd += ["-m", "hmm-tumor"]
                    if "--drop-low-coverage" not in user_options:
                        cmd += ["--drop-low-coverage"]
                # preferentially use conda installed Rscript
                export_cmd = ("%s && export TMPDIR=%s && "
                              % (utils.get_R_exports(), os.path.dirname(tx_out_file)))
                do.run(export_cmd + " ".join(cmd), "CNVkit segment")
    return out_file

def _cnvkit_metrics(cnns, target_bed, antitarget_bed, cov_interval, items):
    """Estimate noise of a sample using a flat background.

    Only used for panel/targeted data due to memory issues with whole genome
    samples.

    """
    if cov_interval == "genome":
        return cnns
    target_cnn = [x["file"] for x in cnns if x["cnntype"] == "target"][0]
    background_file = "%s-flatbackground.cnn" % utils.splitext_plus(target_cnn)[0]
    background_file = cnvkit_background([], background_file, items, target_bed, antitarget_bed)
    cnr_file, data = _cnvkit_fix_base(cnns, background_file, items, "-flatbackground")
    cns_file = _cnvkit_segment(cnr_file, cov_interval, data)
    metrics_file = "%s-metrics.txt" % utils.splitext_plus(target_cnn)[0]
    if not utils.file_exists(metrics_file):
        with file_transaction(data, metrics_file) as tx_metrics_file:
            cmd = [_get_cmd(), "metrics", "-o", tx_metrics_file, "-s", cns_file, "--", cnr_file]
            do.run(_prep_cmd(cmd, tx_metrics_file), "CNVkit metrics")
    metrics = _read_metrics_file(metrics_file)
    out = []
    for cnn in cnns:
        cnn["metrics"] = metrics
        out.append(cnn)
    return out

def _read_metrics_file(in_file):
    with open(in_file) as in_handle:
        header = next(in_handle).strip().split("\t")[1:]
        vals = map(float, next(in_handle).strip().split("\t")[1:])
    return dict(zip(header, vals))

@utils.map_wrap
@zeromq_aware_logging
def _cnvkit_fix(cnns, background_cnn, items, ckouts):
    """Normalize samples, correcting sources of bias.
    """
    return [_cnvkit_fix_base(cnns, background_cnn, items, ckouts)]

def _cnvkit_fix_base(cnns, background_cnn, items, ckouts, ext=""):
    assert len(cnns) == 2, "Expected target and antitarget CNNs: %s" % cnns
    target_cnn = [x["file"] for x in cnns if x["cnntype"] == "target"][0]
    antitarget_cnn = [x["file"] for x in cnns if x["cnntype"] == "antitarget"][0]
    dindex, data = [(i, x) for (i, x) in enumerate(items) if dd.get_sample_name(x) == cnns[0]["sample"]][0]
    out_file = ckouts[dindex]["cnr"]
    return (run_fix(target_cnn, antitarget_cnn, background_cnn, out_file, data), data)

@utils.map_wrap
@zeromq_aware_logging
def run_fix_parallel(target_cnn, antitarget_cnn, background_cnn, out_file, data):
    return [run_fix(target_cnn, antitarget_cnn, background_cnn, out_file, data)]

def run_fix(target_cnn, antitarget_cnn, background_cnn, out_file, data):
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            cmd = [_get_cmd(), "fix", "-o", tx_out_file, target_cnn, antitarget_cnn, background_cnn,
                   "--sample-id", dd.get_sample_name(data)]
            do.run(_prep_cmd(cmd, tx_out_file), "CNVkit fix")
    return out_file

def _select_background_cnns(cnns):
    """Select cnns to use for background calculations.

    Uses background samples in cohort, and will remove CNNs with high
    on target variability. Uses (number of segments * biweight midvariance) as metric
    for variability with higher numbers being more unreliable.
    """
    min_for_variability_analysis = 20
    pct_keep = 0.10
    b_cnns = [x for x in cnns if x["itype"] == "background" and x.get("metrics")]
    assert len(b_cnns) % 2 == 0, "Expect even set of target/antitarget cnns for background"
    if len(b_cnns) >= min_for_variability_analysis:
        b_cnns_w_metrics = []
        for b_cnn in b_cnns:
            unreliability = b_cnn["metrics"]["segments"] * b_cnn["metrics"]["bivar"]
            b_cnns_w_metrics.append((unreliability, b_cnn))
        b_cnns_w_metrics.sort()
        to_keep = int(math.ceil(pct_keep * len(b_cnns) / 2.0) * 2)
        b_cnns = [x[1] for x in b_cnns_w_metrics][:to_keep]
        assert len(b_cnns) % 2 == 0, "Expect even set of target/antitarget cnns for background"
    return [x["file"] for x in b_cnns]

def cnvkit_background(background_cnns, out_file, items, target_bed=None, antitarget_bed=None):
    """Calculate background reference, handling flat case with no normal sample.
    """
    if not utils.file_exists(out_file):
        with file_transaction(items[0], out_file) as tx_out_file:
            cmd = [_get_cmd(), "reference", "-f", dd.get_ref_file(items[0]), "-o", tx_out_file]
            gender = _get_batch_gender(items)
            if gender:
                cmd += ["--sample-sex", gender]
            if len(background_cnns) == 0:
                assert target_bed and antitarget_bed, "Missing CNNs and target BEDs for flat background"
                cmd += ["-t", target_bed, "-a", antitarget_bed]
            else:
                cmd += background_cnns
            do.run(_prep_cmd(cmd, tx_out_file), "CNVkit background")
    return out_file

def _get_batch_gender(items):
    """Retrieve gender for a batch of items if consistent.

    Better not to specify for mixed populations, CNVkit will work
    it out
    https://github.com/bcbio/bcbio-nextgen/commit/1a0e217c8a4d3cee10fa890fb3cfd4db5034281d#r26279752
    """
    genders = set([population.get_gender(x) for x in items])
    if len(genders) == 1:
        gender = genders.pop()
        if gender != "unknown":
            return gender

def targets_w_bins(cnv_file, access_file, target_anti_fn, work_dir, data):
    """Calculate target and anti-target files with pre-determined bins.
    """
    target_file = os.path.join(work_dir, "%s-target.bed" % dd.get_sample_name(data))
    anti_file = os.path.join(work_dir, "%s-antitarget.bed" % dd.get_sample_name(data))
    if not utils.file_exists(target_file):
        target_bin, _ = target_anti_fn()
        with file_transaction(data, target_file) as tx_out_file:
            cmd = [_get_cmd(), "target", cnv_file, "--split", "-o", tx_out_file,
                   "--avg-size", str(target_bin)]
            do.run(_prep_cmd(cmd, tx_out_file), "CNVkit target")
    if not os.path.exists(anti_file):
        _, anti_bin = target_anti_fn()
        with file_transaction(data, anti_file) as tx_out_file:
            # Create access file without targets to avoid overlap
            # antitarget in cnvkit is meant to do this but appears to not always happen
            # after chromosome 1
            tx_access_file = os.path.join(os.path.dirname(tx_out_file), os.path.basename(access_file))
            pybedtools.BedTool(access_file).subtract(cnv_file).saveas(tx_access_file)
            cmd = [_get_cmd(), "antitarget", "-g", tx_access_file, cnv_file, "-o", tx_out_file,
                   "--avg-size", str(anti_bin)]
            do.run(_prep_cmd(cmd, tx_out_file), "CNVkit antitarget")
    return target_file, anti_file

def targets_from_background(back_cnn, work_dir, data):
    """Retrieve target and antitarget BEDs from background CNN file.
    """
    target_file = os.path.join(work_dir, "%s.target.bed" % dd.get_sample_name(data))
    anti_file = os.path.join(work_dir, "%s.antitarget.bed" % dd.get_sample_name(data))
    if not utils.file_exists(target_file):
        with file_transaction(data, target_file) as tx_out_file:
            out_base = tx_out_file.replace(".target.bed", "")
            cmd = [_get_cmd("reference2targets.py"), "-o", out_base, back_cnn]
            do.run(_prep_cmd(cmd, tx_out_file), "CNVkit targets from background")
            shutil.copy(out_base + ".antitarget.bed", anti_file)
    return target_file, anti_file

def _add_seg_to_output(out, data, enumerate_chroms=False):
    """Export outputs to 'seg' format compatible with IGV and GenePattern.
    """
    out_file = "%s.seg" % os.path.splitext(out["cns"])[0]
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            cmd = [os.path.join(os.path.dirname(sys.executable), "cnvkit.py"), "export",
                   "seg"]
            if enumerate_chroms:
                cmd += ["--enumerate-chroms"]
            cmd += ["-o", tx_out_file, out["cns"]]
            do.run(cmd, "CNVkit export seg")
    out["seg"] = out_file
    return out

def _add_cnr_bedgraph_and_bed_to_output(out, data):
    cnr_file = out["cnr"]
    bedgraph_file = os.path.join(_sv_workdir(data), os.path.basename(cnr_file) + ".bedgraph")
    if not utils.file_exists(bedgraph_file):
        with file_transaction(data, bedgraph_file) as tx_out_file:
            cmd = "sed 1d {cnr_file} | cut -f1,2,3,5 > {tx_out_file}"
            do.run(cmd.format(**locals()), "Converting cnr to bedgraph format")
    out["cnr_bedgraph"] = bedgraph_file

    bed_file = os.path.join(_sv_workdir(data), os.path.basename(cnr_file) + ".bed")
    if not utils.file_exists(bed_file):
        with file_transaction(data, bed_file) as tx_out_file:
            cmd = "sed 1d {cnr_file} | cut -f1,2,3,4,5 > {tx_out_file}"
            do.run(cmd.format(**locals()), "Converting cnr to bed format")
    out["cnr_bed"] = bed_file
    return out

def _compatible_small_variants(data, items):
    """Retrieve small variant (SNP, indel) VCFs compatible with CNVkit.
    """
    from bcbio import heterogeneity
    VarFile = collections.namedtuple("VarFile", ["name", "sample", "normal"])
    out = []
    paired = vcfutils.get_paired(items)
    for v in heterogeneity.get_variants(data, include_germline=not paired):
        vrn_file = v["vrn_file"]
        base, ext = utils.splitext_plus(os.path.basename(vrn_file))
        if paired:
            out.append(VarFile(vrn_file, paired.tumor_name, paired.normal_name))
        else:
            out.append(VarFile(vrn_file, dd.get_sample_name(data), None))
    return out

def _add_variantcalls_to_output(out, data, items, is_somatic=False):
    """Call ploidy and convert into VCF and BED representations.
    """
    call_file = "%s-call%s" % os.path.splitext(out["cns"])
    if not utils.file_exists(call_file):
        with file_transaction(data, call_file) as tx_call_file:
            filters = ["--filter", "cn"]
            cmd = [os.path.join(os.path.dirname(sys.executable), "cnvkit.py"), "call"] + \
                  filters + \
                   ["--ploidy", str(ploidy.get_ploidy([data])),
                    "-o", tx_call_file, out["cns"]]
            small_vrn_files = _compatible_small_variants(data, items)
            if len(small_vrn_files) > 0 and _cna_has_values(out["cns"]):
                cmd += ["--vcf", small_vrn_files[0].name, "--sample-id", small_vrn_files[0].sample]
                if small_vrn_files[0].normal:
                    cmd += ["--normal-id", small_vrn_files[0].normal]
                if not is_somatic:
                    cmd += ["-m", "clonal"]
            gender = _get_batch_gender(items)
            if gender:
                cmd += ["--sample-sex", gender]
            do.run(cmd, "CNVkit call ploidy")
    calls = {}
    for outformat in ["bed", "vcf"]:
        out_file = "%s.%s" % (os.path.splitext(call_file)[0], outformat)
        calls[outformat] = out_file
        if not os.path.exists(out_file):
            with file_transaction(data, out_file) as tx_out_file:
                cmd = [os.path.join(os.path.dirname(sys.executable), "cnvkit.py"), "export",
                       outformat, "--sample-id", dd.get_sample_name(data),
                       "--ploidy", str(ploidy.get_ploidy([data])),
                       "-o", tx_out_file, call_file]
                do.run(cmd, "CNVkit export %s" % outformat)
    out["call_file"] = call_file
    out["vrn_bed"] = annotate.add_genes(calls["bed"], data)
    effects_vcf, _ = effects.add_to_vcf(calls["vcf"], data, "snpeff")
    out["vrn_file"] = effects_vcf or calls["vcf"]
    out["vrn_file"] = shared.annotate_with_depth(out["vrn_file"], items)
    return out

def _add_segmetrics_to_output(out, data):
    """Add metrics for measuring reliability of CNV estimates.
    """
    out_file = "%s-segmetrics.txt" % os.path.splitext(out["cns"])[0]
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            cmd = [os.path.join(os.path.dirname(sys.executable), "cnvkit.py"), "segmetrics",
                   "--median", "--iqr", "--ci", "--pi",
                   "-s", out["cns"], "-o", tx_out_file, out["cnr"]]
            # Use less fine grained bootstrapping intervals for whole genome runs
            if dd.get_coverage_interval(data) == "genome":
                cmd += ["--alpha", "0.1", "--bootstrap", "50"]
            else:
                cmd += ["--alpha", "0.01", "--bootstrap", "500"]
            do.run(cmd, "CNVkit segmetrics")
    out["segmetrics"] = out_file
    return out

def _add_gainloss_to_output(out, data):
    """Add gainloss based on genes, helpful for identifying changes in smaller genes.
    """
    out_file = "%s-gainloss.txt" % os.path.splitext(out["cns"])[0]
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            cmd = [os.path.join(os.path.dirname(sys.executable), "cnvkit.py"), "gainloss",
                   "-s", out["cns"], "-o", tx_out_file, out["cnr"]]
            gender = _get_batch_gender([data])
            if gender:
                cmd += ["--sample-sex", gender]
            do.run(cmd, "CNVkit gainloss")
    out["gainloss"] = out_file
    return out

def _add_coverage_bedgraph_to_output(out, data):
    """Add BedGraph representation of coverage to the output
    """
    out_file = "%s.coverage.bedgraph" % os.path.splitext(out["cns"])[0]
    if utils.file_exists(out_file):
        out["bedgraph"] = out_file
        return out
    bam_file = dd.get_align_bam(data)
    bedtools = config_utils.get_program("bedtools", data["config"])
    samtools = config_utils.get_program("samtools", data["config"])
    cns_file = out["cns"]
    bed_file = tempfile.NamedTemporaryFile(suffix=".bed", delete=False).name
    with file_transaction(data, out_file) as tx_out_file:
        cmd = ("sed 1d {cns_file} | cut -f1,2,3 > {bed_file}; "
               "{samtools} view -b -L {bed_file} {bam_file} | "
               "{bedtools} genomecov -bg -ibam - -g {bed_file} >"
               "{tx_out_file}").format(**locals())
        do.run(cmd, "CNVkit bedGraph conversion")
        os.remove(bed_file)
    out["bedgraph"] = out_file
    return out

def _add_plots_to_output(out, data):
    """Add CNVkit plots summarizing called copy number values.
    """
    out["plot"] = {}
    diagram_plot = _add_diagram_plot(out, data)
    if diagram_plot:
        out["plot"]["diagram"] = diagram_plot
    scatter = _add_scatter_plot(out, data)
    if scatter:
        out["plot"]["scatter"] = scatter
    scatter_global = _add_global_scatter_plot(out, data)
    if scatter_global:
        out["plot"]["scatter_global"] = scatter_global
    return out

def _get_larger_chroms(ref_file):
    """Retrieve larger chromosomes, avoiding the smaller ones for plotting.
    """
    from scipy.cluster.vq import kmeans, vq
    all_sizes = []
    for c in ref.file_contigs(ref_file):
        all_sizes.append(float(c.size))
    all_sizes.sort()
    if len(all_sizes) > 5:
        # separate out smaller chromosomes and haplotypes with kmeans
        centroids, _ = kmeans(np.array(all_sizes), 2)
        idx, _ = vq(np.array(all_sizes), centroids)
        little_sizes = tz.first(tz.partitionby(lambda xs: xs[0], zip(idx, all_sizes)))
        little_sizes = [x[1] for x in little_sizes]
        # create one more cluster with the smaller, removing the haplotypes
        centroids2, _ = kmeans(np.array(little_sizes), 2)
        idx2, _ = vq(np.array(little_sizes), centroids2)
        little_sizes2 = tz.first(tz.partitionby(lambda xs: xs[0], zip(idx2, little_sizes)))
        little_sizes2 = [x[1] for x in little_sizes2]
        # get any chromosomes not in haplotype/random bin
        thresh = max(little_sizes2)
    else:
        thresh = 0
    larger_chroms = []
    for c in ref.file_contigs(ref_file):
        if c.size > thresh:
            larger_chroms.append(c.name)
    return larger_chroms

def _remove_haplotype_chroms(in_file, data):
    """Remove shorter haplotype chromosomes from cns/cnr files for plotting.
    """
    larger_chroms = set(_get_larger_chroms(dd.get_ref_file(data)))
    out_file = os.path.join(_sv_workdir(data), "%s-chromfilter%s" % utils.splitext_plus(os.path.basename(in_file)))
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            with open(in_file) as in_handle:
                with open(tx_out_file, "w") as out_handle:
                    for line in in_handle:
                        if ((line.find("chromosome") >= 0 and line.find("start") >= 0)
                              or line.split()[0] in larger_chroms):
                            out_handle.write(line)
    return out_file

def _add_global_scatter_plot(out, data):
    out_file = os.path.join(_sv_workdir(data),
                            os.path.splitext(os.path.basename(out["cnr"]))[0] + "-scatter_global.pdf")
    if utils.file_exists(out_file):
        return out_file
    cnr = _remove_haplotype_chroms(out["cnr"], data)
    cns = _remove_haplotype_chroms(out["cns"], data)
    with file_transaction(data, out_file) as tx_out_file:
        cmd = [_get_cmd(), "scatter", "-s", cns, "-o", tx_out_file, cnr]
        do.run(_prep_cmd(cmd, tx_out_file), "CNVkit global scatter plot")
    return out_file

def _add_scatter_plot(out, data):
    out_file = os.path.join(_sv_workdir(data),
                            os.path.splitext(os.path.basename(out["cnr"]))[0] + "-scatter.pdf")
    priority_bed = dd.get_svprioritize(data)
    if not priority_bed:
        return None
    priority_bed = plot._prioritize_plot_regions(pybedtools.BedTool(priority_bed), data, os.path.dirname(out_file))
    if utils.file_exists(out_file):
        return out_file
    cnr = _remove_haplotype_chroms(out["cnr"], data)
    cns = _remove_haplotype_chroms(out["cns"], data)
    with file_transaction(data, out_file) as tx_out_file:
        cmd = [_get_cmd(), "scatter", "-s", cns, "-o", tx_out_file, "-l",
               priority_bed, cnr]
        do.run(_prep_cmd(cmd, tx_out_file), "CNVkit scatter plot")
    return out_file

def _cnx_is_empty(in_file):
    """Check if cnr or cns files are empty (only have a header)
    """
    with open(in_file) as in_handle:
        for i, line in enumerate(in_handle):
            if i > 0:
                return False
    return True

def _add_diagram_plot(out, data):
    out_file = os.path.join(_sv_workdir(data),
                            os.path.splitext(os.path.basename(out["cnr"]))[0] + "-diagram.pdf")
    cnr = _remove_haplotype_chroms(out["cnr"], data)
    cns = _remove_haplotype_chroms(out["cns"], data)
    if _cnx_is_empty(cnr) or _cnx_is_empty(cns):
        return None
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            cmd = [_get_cmd(), "diagram", "-s", cns,
                   "-o", tx_out_file, cnr]
            gender = _get_batch_gender([data])
            if gender:
                cmd += ["--sample-sex", gender]
            do.run(_prep_cmd(cmd, tx_out_file), "CNVkit diagram plot")
    return out_file

def segment_from_cnr(cnr_file, data, out_base):
    """Provide segmentation on a cnr file, used in external PureCN integration.
    """
    cns_file = _cnvkit_segment(cnr_file, dd.get_coverage_interval(data),
                               data, [data], out_file="%s.cns" % out_base, detailed=True)
    out = _add_seg_to_output({"cns": cns_file}, data, enumerate_chroms=False)
    return out["seg"]

# ## Theta support

def export_theta(ckout, data):
    """Provide updated set of data with export information for TheTA2 input.
    """
    cns_file = chromhacks.bed_to_standardonly(ckout["cns"], data, headers="chromosome")
    cnr_file = chromhacks.bed_to_standardonly(ckout["cnr"], data, headers="chromosome")
    out_file = "%s-theta.input" % utils.splitext_plus(cns_file)[0]
    if not utils.file_exists(out_file):
        with file_transaction(data, out_file) as tx_out_file:
            cmd = [_get_cmd(), "export", "theta", cns_file, cnr_file, "-o", tx_out_file]
            do.run(_prep_cmd(cmd, tx_out_file), "Export CNVkit calls as inputs for TheTA2")
    ckout["theta_input"] = out_file
    return ckout


# -*- coding: utf-8 -*-
#
# Developed by Chaoyu Huang
# email: b608390.cs04@nctu.edu.tw
# Lot's of code reference to Pix2Vox:
# https://github.com/hzxie/Pix2Vox
#

import json
import numpy as np
import os
import torch
import torch.backends.cudnn
import torch.utils.data

import utils.point_cloud_visualization
import utils.data_loaders
import utils.data_transforms
import utils.network_utils
import utils.view_pred_utils
import utils.rotation_eval

from datetime import datetime as dt

from models.encoder import Encoder
from models.decoder import Decoder
from models.view_estimater import ViewEstimater

from losses.chamfer_loss import ChamferLoss
from losses.earth_mover_distance import EMD
from losses.chamfer_loss import ChamferLoss
from losses.cross_entropy_loss import CELoss
from losses.delta_loss import DeltaLoss

def test_net(cfg):
    # Enable the inbuilt cudnn auto-tuner to find the best algorithm to use
    torch.backends.cudnn.benchmark = True

    # Set up data augmentation
    IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W
    CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W

    test_transforms = utils.data_transforms.Compose([
        utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE),
        utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE),
        utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD),
        utils.data_transforms.ToTensor(),
    ])
    dataset_loader = utils.data_loaders.DATASET_LOADER_MAPPING[cfg.DATASET.TEST_DATASET](cfg)
    test_data_loader = torch.utils.data.DataLoader(dataset=dataset_loader.get_dataset(
                                                   utils.data_loaders.DatasetType.TEST,  test_transforms),
                                                   batch_size=cfg.CONST.BATCH_SIZE,
                                                   num_workers=1,
                                                   pin_memory=True,
                                                   shuffle=False)

    # Set up networks
    encoder = Encoder(cfg)
    decoder = Decoder(cfg)
    azi_classes, ele_classes = int(360 / cfg.CONST.BIN_SIZE), int(180 / cfg.CONST.BIN_SIZE)
    view_estimater = ViewEstimater(cfg, azi_classes=azi_classes, ele_classes=ele_classes)

    if torch.cuda.is_available():
        encoder = torch.nn.DataParallel(encoder).cuda()
        decoder = torch.nn.DataParallel(decoder).cuda()
        view_estimater = torch.nn.DataParallel(view_estimater).cuda()

    # Load weight
    # Load weight for encoder, decoder
    print('[INFO] %s Loading reconstruction weights from %s ...' % (dt.now(), cfg.TEST.RECONSTRUCTION_WEIGHTS))
    rec_checkpoint = torch.load(cfg.TEST.RECONSTRUCTION_WEIGHTS)
    encoder.load_state_dict(rec_checkpoint['encoder_state_dict'])
    decoder.load_state_dict(rec_checkpoint['decoder_state_dict'])
    print('[INFO] Best reconstruction result at epoch %d ...' % rec_checkpoint['epoch_idx'])

    # Load weight for view estimater
    print('[INFO] %s Loading view estimation weights from %s ...' % (dt.now(), cfg.TEST.VIEW_ESTIMATION_WEIGHTS))
    view_checkpoint = torch.load(cfg.TEST.VIEW_ESTIMATION_WEIGHTS)
    view_estimater.load_state_dict(view_checkpoint['view_estimator_state_dict'])
    print('[INFO] Best view estimation result at epoch %d ...' % view_checkpoint['epoch_idx'])

    # Set up loss functions
    emd = EMD().cuda()
    cd = ChamferLoss().cuda()

    # Batch average meterics
    cd_distances = utils.network_utils.AverageMeter()
    emd_distances = utils.network_utils.AverageMeter()
    pointwise_emd_distances = utils.network_utils.AverageMeter()

    test_preds = torch.zeros([1, 3], dtype=torch.float).cuda()
    test_ground_truth_views = torch.zeros([1, 3], dtype=torch.long).cuda()

    # Switch models to evaluation mode
    encoder.eval()
    decoder.eval()
    view_estimater.eval()

    n_batches = len(test_data_loader)

    # Testing loop
    for sample_idx, (taxonomy_names, sample_names, rendering_images,
                    ground_truth_point_clouds, ground_truth_views) in enumerate(test_data_loader):
        with torch.no_grad():
            # Only one image per sample
            rendering_images = torch.squeeze(rendering_images, 1)

             # Get data from data loader
            rendering_images = utils.network_utils.var_or_cuda(rendering_images)
            ground_truth_point_clouds = utils.network_utils.var_or_cuda(ground_truth_point_clouds)
            ground_truth_views = utils.network_utils.var_or_cuda(ground_truth_views)

            #=================================================#
            #           Test the encoder, decoder             #
            #=================================================#
            vgg_features, image_code = encoder(rendering_images)
            image_code = [image_code]
            generated_point_clouds = decoder(image_code)

            # Compute CD, EMD
            cd_distance = cd(generated_point_clouds, ground_truth_point_clouds) / cfg.CONST.BATCH_SIZE / cfg.CONST.NUM_POINTS
            emd_distance = torch.mean(emd(generated_point_clouds, ground_truth_point_clouds))
            pointwise_emd_distance = emd_distance / cfg.CONST.NUM_POINTS

            #=================================================#
            #          Test the view estimater                #
            #=================================================#
            # output[0]:prediction of azi class
            # output[1]:prediction of ele class
            # output[2]:prediction of azi regression
            # output[3]:prediction of ele regression
            output = view_estimater(vgg_features)

            #=================================================#
            #              Get predict view                   #
            #=================================================#
            preds_cls = utils.view_pred_utils.get_pred_from_cls_output([output[0], output[1]])

            preds = []
            for n in range(len(preds_cls)):
                pred_delta = output[n + 2]
                delta_value = pred_delta[torch.arange(pred_delta.size(0)), preds_cls[n].long()].tanh() / 2
                preds.append((preds_cls[n].float() + delta_value + 0.5) * cfg.CONST.BIN_SIZE)
            # In this experiment, expect all the object has a fixed rotaion angle 0 and with diferent view point
            # Add the zero inplane term for the rotation acc calculation
            zero_inplane = torch.zeros_like(preds[0])
            # Add fixed inplane rotation to  pred view
            test_pred = torch.cat((preds[0].unsqueeze(1), preds[1].unsqueeze(1), zero_inplane.unsqueeze(1)), 1)
            # Add fixed inplane rotation to ground_truth_views
            zero_inplane = zero_inplane.unsqueeze(1)
            ground_truth_views = torch.cat((ground_truth_views, zero_inplane.long()), 1)

            # Append loss and accuracy to average metrics
            cd_distances.update(cd_distance.item())
            emd_distances.update(emd_distance.item())
            pointwise_emd_distances.update(pointwise_emd_distance.item())

            # concatenate results and labels for view estimation
            test_preds = torch.cat((test_preds, test_pred), 0)
            test_ground_truth_views = torch.cat((test_ground_truth_views, ground_truth_views), 0)

            print("Test on [%d/%d] data, CD: %.4f EMD: %.4f Total EMD %.4f" % (sample_idx + 1,  n_batches, cd_distance.item(), pointwise_emd_distance.item(), emd_distance.item()))

    test_preds = test_preds[1:, :]
    test_ground_truth_views = test_ground_truth_views[1:, :]

    # calculate the rotation errors between prediction and ground truth
    test_errs = utils.rotation_eval.rotation_err(test_preds, test_ground_truth_views.float()).cpu().numpy()
    Acc = 100. * np.mean(test_errs <= 30)
    Med = np.median(test_errs)

   # print result
    print("Reconstruction result:")
    print("CD result: ", cd_distances.avg)
    print("Pointwise EMD result: ", pointwise_emd_distances.avg)
    print("Total EMD result", emd_distances.avg)
    print("View estimation result:")
    print('Med_Err is %.2f, and Acc_pi/6 is %.2f \n \n' % (Med, Acc))
    logname = cfg.TEST.RESULT_PATH
    with open(logname, 'a') as f:
        f.write('Reconstruction result: \n')
        f.write("CD result: %.8f \n" % cd_distances.avg)
        f.write("Pointwise EMD result: %.8f \n" % pointwise_emd_distances.avg)
        f.write("Total EMD result: %.8f \n" % emd_distances.avg)
        f.write('View estimation result: \n')
        f.write('Med_Err is %.2f, and Acc_pi/6 is %.2f \n \n' % (Med, Acc))


import unittest
import shutil
import numpy as np
import pandas as pd
from sklearn import datasets
from supervised import AutoML

class EnsembleSaveLoadTest(unittest.TestCase):

    automl_dir = "automl_1"

    def tearDown(self):
        shutil.rmtree(self.automl_dir, ignore_errors=True)

    def test_save_load(self):
        a = AutoML(
            results_path=self.automl_dir,
            total_time_limit=10,
            explain_level=0,
            mode="Explain",
            train_ensemble=True,
        )
        a.set_advanced(start_random_models=1)

        X, y = datasets.make_classification(
            n_samples=100,
            n_features=5,
            n_informative=4,
            n_redundant=1,
            n_classes=2,
            n_clusters_per_class=3,
            n_repeated=0,
            shuffle=False,
            random_state=0,
        )
        X = pd.DataFrame(X, columns=[f"f_{i}" for i in range(X.shape[1])])

        a.fit(X, y)
        p = a.predict(X)

        a2 = AutoML(results_path=self.automl_dir)
        p2 = a2.predict(X)

        self.assertTrue((p["label"] == p2["label"]).all())


"""
This file contains a function to decode a single synthetic tree.
"""
import pandas as pd
import numpy as np
from syn_net.utils.data_utils import ReactionSet
from syn_net.utils.predict_utils import synthetic_tree_decoder_multireactant, load_modules_from_checkpoint, mol_fp

# define some constants (here, for the Hartenfeller-Button test set)
nbits        = 4096
out_dim      = 256
rxn_template = 'hb'
featurize    = 'fp'
param_dir    = 'hb_fp_2_4096_256'
ncpu         = 1

# load the purchasable building block embeddings
bb_emb = np.load('/pool001/whgao/data/synth_net/st_hb/enamine_us_emb_fp_256.npy')

# define path to the reaction templates and purchasable building blocks
path_to_reaction_file   = f'/pool001/whgao/data/synth_net/st_{rxn_template}/reactions_{rxn_template}.json.gz'
path_to_building_blocks = f'/pool001/whgao/data/synth_net/st_{rxn_template}/enamine_us_matched.csv.gz'

# define paths to pretrained modules
param_path  = f'/home/whgao/synth_net/synth_net/params/{param_dir}/'
path_to_act = f'{param_path}act.ckpt'
path_to_rt1 = f'{param_path}rt1.ckpt'
path_to_rxn = f'{param_path}rxn.ckpt'
path_to_rt2 = f'{param_path}rt2.ckpt'

# load the purchasable building block SMILES to a dictionary
building_blocks = pd.read_csv(path_to_building_blocks, compression='gzip')['SMILES'].tolist()
bb_dict         = {building_blocks[i]: i for i in range(len(building_blocks))}

# load the reaction templates as a ReactionSet object
rxn_set = ReactionSet()
rxn_set.load(path_to_reaction_file)
rxns    = rxn_set.rxns

# load the pre-trained modules
act_net, rt1_net, rxn_net, rt2_net = load_modules_from_checkpoint(
    path_to_act=path_to_act,
    path_to_rt1=path_to_rt1,
    path_to_rxn=path_to_rxn,
    path_to_rt2=path_to_rt2,
    featurize=featurize,
    rxn_template=rxn_template,
    out_dim=out_dim,
    nbits=nbits,
    ncpu=ncpu,
)

def func(smi):
    """
    Generates the synthetic tree for the input molecular embedding.

    Args:
        smi (str): SMILES string corresponding to the molecule to decode.

    Returns:
        smi (str): SMILES for the final chemical node in the tree.
        similarity (float): Similarity measure between the final chemical node
            and the input molecule.
        tree (SyntheticTree): The generated synthetic tree.
    """
    emb = mol_fp(smi)
    try:
        smi, similarity, tree, action = synthetic_tree_decoder_multireactant(
            z_target=emb,
            building_blocks=building_blocks,
            bb_dict=bb_dict,
            reaction_templates=rxns,
            mol_embedder=mol_fp,
            action_net=act_net,
            reactant1_net=rt1_net,
            rxn_net=rxn_net,
            reactant2_net=rt2_net,
            bb_emb=bb_emb,
            rxn_template=rxn_template,
            n_bits=nbits,
            beam_width=3,
            max_step=15)
    except Exception as e:
        print(e)
        action = -1

    if action != 3:
        return None, 0, None
    else:
        return smi, similarity, tree