#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Feb 22 10:50:35 2021

@author: oskar.hickl
"""

import logging
import sys
import os
import glob
import subprocess
import argparse

print("=== PYTHON EXECUTABLE ===")
print(sys.executable)
print("=== PYTHON VERSION ===")
print(sys.version)

# required_packages = [
#     "numpy",
#     "pandas",
#     "matplotlib",
#     "networkx",
#     "seaborn",
#     "joblib",
#     "scikit-learn",
#     "scikit-bio",
#     "openTSNE",
#     "hdbscan",
# ]

# for pkg in required_packages:
#     subprocess.run(f"python -m pip install --no-cache-dir {pkg}", shell=True, check=True)

# binny_out = snakemake.params['binny_out']
# sample = snakemake.params['sample']
# mg_depth_file = snakemake.input['mgdepth']
# assembly = snakemake.input['assembly']
# annot_file = snakemake.params['gff']
# raw_annot = snakemake.input['raw_gff']
# tigrfam2pfam_file = snakemake.params['t2p']
# taxon_marker_set_file = snakemake.params['marker_sets']
# prokka_checkm_marker_hmm_out = snakemake.input['hmm_markers']
# functions = snakemake.params['py_functions']
# min_purity = float(snakemake.params['purity'])
# min_completeness = float(snakemake.params['min_completeness'])
# starting_completeness = float(snakemake.params['start_completeness'])
# kmers = snakemake.params['kmers']
# mask_disruptive_sequences = eval(snakemake.params['mask_disruptive_sequences'])
# extract_scmags = eval(snakemake.params['extract_scmags'])
# coassembly_mode = snakemake.params['coassembly_mode']
# nx_val = int(snakemake.params['nx_val'])
# min_contig_length = int(snakemake.params['min_cutoff'])
# max_contig_length = int(snakemake.params['max_cutoff'])
# min_contig_length_marker = int(snakemake.params['min_cutoff_marker'])
# max_contig_length_marker = int(snakemake.params['max_cutoff_marker'])
# max_contig_threshold = float(snakemake.params['max_n_contigs'])
# max_marker_lineage_depth_lvl = int(snakemake.params['max_marker_lineage_depth_lvl'])
# max_embedding_tries = int(snakemake.params['max_embedding_tries'])
# include_depth_initial = eval(snakemake.params['include_depth_initial'])
# include_depth_main = eval(snakemake.params['include_depth_main'])
# hdbscan_epsilon_range = [float(epsilon) for epsilon in snakemake.params['hdbscan_epsilon_range'].split(',')]
# hdbscan_min_samples_range = [int(min_sample) for min_sample in snakemake.params['hdbscan_min_samples_range'].split(',')]
# dist_metric = snakemake.params['distance_metric']
# write_contig_data = eval(snakemake.params['write_contig_data'])

parser = argparse.ArgumentParser()
parser.add_argument("--binny-out", required=True)
parser.add_argument("--sample", required=True)
parser.add_argument("--mgdepth", required=True)
parser.add_argument("--assembly", required=True)
parser.add_argument("--gff", required=True)
parser.add_argument("--raw-gff", required=True)
parser.add_argument("--t2p", required=True)
parser.add_argument("--marker-sets", required=True)
parser.add_argument("--hmm-markers", required=True)
parser.add_argument("--py-functions", required=True)
parser.add_argument("--purity", required=True)
parser.add_argument("--min-completeness", required=True)
parser.add_argument("--start-completeness", required=True)
parser.add_argument("--kmers", required=True)
parser.add_argument("--mask-disruptive-sequences", required=True, type=bool)
parser.add_argument("--extract-scmags", required=True, type=bool)
parser.add_argument("--coassembly-mode", required=True)
parser.add_argument("--min-cutoff", required=True)
parser.add_argument("--max-cutoff", required=True)
parser.add_argument("--min-cutoff-marker", required=True)
parser.add_argument("--max-cutoff-marker", required=True)
parser.add_argument("--nx-val", required=True)
parser.add_argument("--max-n-contigs", required=True)
parser.add_argument("--max-marker-lineage-depth-lvl", required=True)
parser.add_argument("--max-embedding-tries", required=True)
parser.add_argument("--include-depth-initial", required=True, type=bool)
parser.add_argument("--include-depth-main", required=True, type=bool)
parser.add_argument("--hdbscan-epsilon-range", required=True)
parser.add_argument("--hdbscan-min-samples-range", required=True)
parser.add_argument("--distance-metric", required=True)
parser.add_argument("--write-contig-data", required=True, type=bool)
parser.add_argument("--threads", required=True, type=int)
parser.add_argument("--log", required=True)
args = parser.parse_args()

binny_out = args.binny_out
sample = args.sample
mg_depth_file = args.mgdepth
assembly = args.assembly
annot_file = args.gff
raw_annot = args.raw_gff
tigrfam2pfam_file = args.t2p
taxon_marker_set_file = args.marker_sets
prokka_checkm_marker_hmm_out = args.hmm_markers
functions = args.py_functions
min_purity = float(args.purity)
min_completeness = float(args.min_completeness)
starting_completeness = float(args.start_completeness)
kmers = args.kmers
mask_disruptive_sequences = args.mask_disruptive_sequences
extract_scmags = args.extract_scmags
coassembly_mode = args.coassembly_mode
nx_val = int(args.nx_val)
min_contig_length = int(args.min_cutoff)
max_contig_length = int(args.max_cutoff)
min_contig_length_marker = int(args.min_cutoff_marker)
max_contig_length_marker = int(args.max_cutoff_marker)
max_contig_threshold = float(args.max_n_contigs)
max_marker_lineage_depth_lvl = int(args.max_marker_lineage_depth_lvl)
max_embedding_tries = int(args.max_embedding_tries)
include_depth_initial = args.include_depth_initial
include_depth_main = args.include_depth_main
hdbscan_epsilon_range = [float(epsilon) for epsilon in args.hdbscan_epsilon_range.split(',')]
hdbscan_min_samples_range = [int(min_sample) for min_sample in args.hdbscan_min_samples_range.split(',')]
dist_metric = args.distance_metric
write_contig_data = args.write_contig_data

intermediary_file_dir = 'intermediary'

threads = args.threads
log = args.log

n_dim = 2

sys.path.append(os.path.dirname(functions))
from binny_functions import *

# To achieve reproducible results with HDBSCAN and ensure same seed, because other tools that accept seed arguments,
# might mess with the global numpy seed
np.random.seed(0)

logging.basicConfig(filename=log, level=logging.INFO, format='%(asctime)s - %(message)s',  # logging.INFO
                    datefmt='%d/%m/%Y %I:%M:%S %p', filemode='w')

numba_logger = logging.getLogger('numba')
numba_logger.setLevel(logging.WARNING)

logging.info('Starting Binny run for sample {0}.'.format(sample))


# Check if bin dir empty
if glob.glob(os.path.join(binny_out, "bins/*.fasta")):
    print(os.listdir(os.path.join(binny_out, "bins/")), file=sys.stderr)
    print('Bin dir contains fasta files. Move or delete. Exiting.')
    raise Exception

# Load TIGRFAMs to PFAMs conversion table.
tigrfam2pfam_data = tigrfam2pfam_dict(tigrfam2pfam_file)

# Merge Prokka gff with marker set data, load annotation df, and load assembly.
checkm_hmmer_search2prokka_gff(prokka_checkm_marker_hmm_out, raw_annot, gff_out_path=os.path.join(binny_out,
                                                                                                  'intermediary'))
annot_df, annot_dict = gff2ess_gene_df(annot_file, target_attribute='checkm_marker', get_dict=True)
assembly_dict = load_fasta(assembly)

# Build marker set graph db.
taxon_marker_sets = load_checkm_markers(taxon_marker_set_file)

# Look for complete genomes on single contigs
all_good_bins = {}
if extract_scmags:
    logging.info('Looking for single contig bins.')
    single_contig_bins = get_single_contig_bins(annot_df, all_good_bins, n_dim, taxon_marker_sets, tigrfam2pfam_data,
                                                threads, max_marker_lineage_depth_lvl=max_marker_lineage_depth_lvl)
    logging.info('Found {0} single contig bins.'.format(len(single_contig_bins)))
else:
    single_contig_bins = []

# logging.info('Getting assembly dict without scMAG.'.format(len(single_contig_bins)))
# assembly_dict_wo_scmags = {key: val for key, val in assembly_dict.items() if key in
#                            set(assembly_dict.keys()).difference(set(single_contig_bins))}

logging.info(f'Calculating N{nx_val}'.format(len(single_contig_bins)))
nx = calc_assembly_nx(assembly_dict, single_contig_bins, nx_val)
nx2 = calc_assembly_nx(assembly_dict, [], nx_val)
logging.info(f'N{nx_val} is {nx}, with scMAGs would be {nx2}.'.format(len(single_contig_bins)))
min_contig_length = min(max(nx, min_contig_length), max_contig_length)
min_contig_length_marker = min(max(int(nx / 3), min_contig_length_marker), max_contig_length_marker)

# Load assembly and mask rRNAs and CRISPR arrays
# contig_list = [[contig] + [seq] for contig, seq in assembly_dict.items() if (len(seq) >= min_contig_length
#                                                                              or (annot_dict.get(contig)
#                                                                              and len(seq) >= min_contig_length_marker))
#                                                                             and contig not in single_contig_bins]
contig_list = [[contig] + [seq] for contig, seq in assembly_dict.items() if contig not in single_contig_bins
                                                                            and len(seq) >= 500]

# logging.info('{0} contigs match length threshold of {1}bp or contain marker genes and'
#              ' have a size of at least {2}bp'.format(len(contig_list), min_contig_length, min_contig_length_marker))

contig_rrna_crispr_region_dict = gff2low_comp_feature_dict(annot_file)
if mask_disruptive_sequences:
    logging.info('Masking potentially disruptive sequences from k-mer counting.')
    mask_rep_featrues(contig_rrna_crispr_region_dict, contig_list)
else:
    logging.info('Not masking potentially disruptive sequences from k-mer counting.')

# Get length normalized k-mer frequencies.
kmer_sizes = [int(kmer) for kmer in kmers.split(',')]

start = timer()
kfreq_array = get_contig_kmer_matrix(contig_list, kmer_sizes, threads)
end = timer()
logging.info('K-mer frequency matrix created in {0}s.'.format(int(end - start)))

# Make array, removing fully masked sequences with no counts and standardize k-mer freq data
x = np.array([c_kfreq[1:] for c_kfreq in kfreq_array[1:] if not sum(c_kfreq[1:]) == 0])
x_contigs = [c_kfreq[0] for c_kfreq in kfreq_array[1:] if not sum(c_kfreq[1:]) == 0]

main_contig_data_dict = {cont: seq for cont, seq in zip(x_contigs, x)}

# Load depth data
depth_dict = load_depth_dict(mg_depth_file)
n_depth_samples = list(depth_dict.values())[0].shape[0]
logging.debug(f'list(depth_dict.values())[0].shape[0]: {n_depth_samples}')

if coassembly_mode == 'on' or (coassembly_mode == 'auto' and n_depth_samples > 1):
    min_contig_length_marker = 500
    if n_depth_samples >= 3:
        include_depth_main = True
    logging.info('Using coassembly mode.')
else:
    logging.info('Using single sample mode.')

# Run iterative dimension reduction, manifold learning, cluster detection and assessment.
all_good_bins, contig_data_df_org, min_purity = iterative_embedding(x_contigs, depth_dict, all_good_bins, starting_completeness,
                                                        min_purity, min_completeness, threads, n_dim, annot_file,
                                                        mg_depth_file, single_contig_bins, taxon_marker_sets,
                                                        tigrfam2pfam_data, main_contig_data_dict, assembly_dict,
                                                        max_contig_threshold, min_contig_length_marker,
                                                        include_depth_initial, max_embedding_tries,
                                                        include_depth_main, hdbscan_epsilon_range,
                                                        hdbscan_min_samples_range, dist_metric,
                                                        contigs2clusters_out_path=os.path.join(binny_out, 'intermediary'),
                                                        max_marker_lineage_depth_lvl=max_marker_lineage_depth_lvl,
                                                        coassembly_mode=coassembly_mode)

all_contigs = []
for bin in all_good_bins:
    all_contigs.extend(all_good_bins[bin]['contigs'])
if len(all_contigs) != len(set(all_contigs)):
    logging.warning('WARNING: {0} duplicate contigs in bins found!'.format(len(all_contigs) - len(set(all_contigs))))

# Write bin fastas.
write_bins(all_good_bins, assembly, min_comp=int(min_completeness), min_pur=int(min_purity),
           bin_dir=os.path.join(binny_out, 'bins'))

bin_dict = {contig: bin for bin, bin_data in all_good_bins.items() for contig in bin_data['contigs']}

all_cont_data_dict = {}

for contig, k_freq in main_contig_data_dict.items():
    all_cont_data_dict[contig] = {'bin': bin_dict.get(contig, 'N'),
                                  'k-mer_freqs': ';'.join([str(k) for k in list(k_freq)]),
                                  'depths': ';'.join([str(d) for d in list(depth_dict.get(contig))])}

all_cont_data_dict.update({contig: {'bin': contig,
                                  'k-mer_freqs': '',
                                  'depths': ';'.join([str(d) for d in list(depth_dict.get(contig))])}
                           for contig in single_contig_bins})

contig_data_df = pd.DataFrame.from_dict(all_cont_data_dict, orient='index', columns=['bin', 'k-mer_freqs', 'depths'])
# Downcast df
downcast_pd_df(contig_data_df)

if write_contig_data:
    logging.info('Writing contig data to file.')
    compression_opts = dict(method='gzip')
    contig_data_df.to_csv(os.path.join(binny_out, 'contig_data.tsv.gz'), header=True, index=True, index_label='contig',
                          chunksize=250000, compression=compression_opts, sep='\t')

os.mknod(os.path.join(binny_out, "binny.done"))

logging.info('Run finished.')
