#!/usr/bin/env python2.7
#*-* coding: utf-8 *-*

"""
Load the MANTA2 database from data files generated by the remap process
performed in Anthony Mathalier's lab. The imput files/formats have changed
since the original MANTA. The main difference is we no longer have or care
to capture the peak information.

The master CSV file has a new format:
    FILE_NAME, ENCODE_ID, CELL_LINE, TF, JASPAR_TF, JASPAR_ID, VERSION
    e.g. ENCSR000AHD.ctcf.mcf7,ENCSR000AHD,MCF7,CTCF,CTCF,MA0139,1

We have a BED file for each TF.

These BED files are named:

   <jaspar_matrix_id>_<tf_name>.bed    e.g. MA0651.1_HOXC11.bed

Each line of each of these BED files relates which experiment(s) contain this
specific binding site for the given TF.

The files are tab delimited with the fields:
    chrom, start, end, TF name, rel. score, abs. score, strand, exp_info;exp_info;...

    e.g. chr1       629919          629930       HOXC11  7.831;86.3      +       GSE54027.hoxc11.mcf7_ly2_tamoxifen;GSE54027.hoxc11.mcf7_ly2

From the top directory there are subdirectories for each TF (named by the
JASPAR matrix ID). These contain text files names:

    <jaspar_matrix_id>_<tf_name>_SNV_impacts.txt
    e.g. MA0651.1/MA0651.1_HOXC11_SNV_impacts.txt
"""

import os
import sys
import argparse
import pymongo
import bson
import logging
import datetime
import fnmatch

from pymongo import MongoClient

MANTA2_DB_HOST = 'manta.cmmt.ubc.ca'
MANTA2_DB_NAME = 'manta2'
MANTA2_DB_USER = 'manta2_rw'

EXPERIMENT_CSV_FILENAME = 'mapping_to_JASPAR.csv'

LOGFILE = 'load_manta2_db.log'
REMAP_DATA_TYPE = 'ChIP-seq'    # in future may have different types
SPECIES = 9606

DATA_SOURCE_NAMES = {
    'ENC' : 'ENCODE',
    'GSE' : 'GEO',
    'ERP' : 'ArrayExpress',
    'ERR' : 'ArrayExpress'
}

DATA_SOURCE_URLS = {
    'ENCODE' : 'https://www.encodeproject.org/experiments/<experiment_id>',
    'ArrayExpress' : 'https://www.ebi.ac.uk/arrayexpress/search.html?query=<experiment_id>',
    'GEO' : 'https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi?acc=<experiment_id>'
}

# Base path for the TFBS data files.
TFBS_BASE_PATH = '/raid9/dave/MANTA2/data/20170929_variant_scoring'


def clean_database(db):
    logging.info("Cleaning database\n")

    # Drop the collections
    logging.info("Dropping experiments collection\n")
    db.experiments.drop()
    logging.info("Dropping tfbs_snvs collection\n")
    db.tfbs_snvs.drop()

    # Re-create the indexes
    #logging.info("Creating experiments indexes\n")
    #db.experiments.ensure_index(['experiment_id', 'cell_line', 'tf_name'])

    #logging.info("Creating tfbs_snvs indexes\n")
    #db.tfbs_snvs.ensure_index([
    #                        ('matrix_id', pymongo.ASCENDING),
    #                        ('chrom', pymongo.ASCENDING),
    #                        ('start', pymongo.ASCENDING),
    #                        ('strand', pymongo.ASCENDING)])

    #db.tfbs_snvs.ensure_index([
    #                        ('chrom', pymongo.ASCENDING),
    #                        ('snvs.pos', pymongo.ASCENDING)])


def load_experiments(db, dir):
    """Load all experiments from the master CSV file
    """

    exp_file = os.path.join(dir, EXPERIMENT_CSV_FILENAME)

    with open(exp_file, 'r') as xf:
        for line in xf:
            line = line.rstrip()

            (experiment_id, data_source_name, cell_line, tf_name,
             jaspar_tf_name, matrix_base_id, matrix_version) = line.split(',')
            

            # Skip header
            if experiment_id == "EXPERIMENT_ID":
                continue

            jaspar_matrix_id = "{}.{}".format(matrix_base_id, matrix_version)

            #data_source_abbrev = experiment_id[0:3]
            #data_source_name = DATA_SOURCE_NAMES[data_source_abbrev]

            data_source_url = DATA_SOURCE_URLS[data_source_name]

            data_source_url = data_source_url.replace(
                                "<experiment_id>", experiment_id)

            # Check if there is already an existing experiment for this TF
            curs = db.experiments.find(
                {
                    "experiment_id" : experiment_id,
                    "cell_line"     : cell_line,
                    "tf_name"       : tf_name
                }
            )

            if not curs or curs.count() == 0:
                # New experiment
                experiment = {
                    "data_source"     : data_source_name,
                    "species"         : SPECIES,
                    "experiment_id"   : experiment_id,
                    "cell_line"       : cell_line,
                    "url"             : data_source_url,
                    "tf_name"         : tf_name,
                    "jaspar_tfs"      : [{"matrix_id"  : jaspar_matrix_id,
                                         "name" : jaspar_tf_name}]
                }
            elif curs.count() > 1:
                    errstr = "More than one experiment record found which matches {}:{}:{}".format(experiment_id, cell_line, tf_name)
                    logging.error(errstr)
                    sys.exit(errstr)
            else:
                # XXX If we were being super careful here, we should also
                # check that this JASPAR TF is not already stored in the
                # list of JASPAR TFs. XXX
                experiment = curs.next()

                new_jaspar_tf = {"matrix_id" : jaspar_matrix_id,
                                 "name"      : jaspar_tf_name}

                if new_jaspar_tf not in experiment['jaspar_tfs']:
                    experiment['jaspar_tfs'].append(new_jaspar_tf)

            # With the "upsert" (third parameter) set to True, if the
            # experiment does not exist already then a new record is inserted,
            # otherwise the existing record is updated.
            db.experiments.update(
                {
                    "experiment_id" : experiment_id,
                    "cell_line"     : cell_line,
                    "tf_name"       : tf_name
                }, experiment, True)


def load_tfbs_snv(db, dir, bed_file=None):
    """Load the TFBS/SNV records.
    The base directory contains subdirectories for each JASPAR motif
    which use the matrix ID as the sub dir name, as well as BED files
    containing the experiment to TFBS mapping files.
    
    These BED files contain the unaltered TFBS positions and scores
    and their association to the specific experiments which were performed
    for the given TF, i.e. have peaks which where scanned for TFBS of the
    JASPAR motif assocated to this TF. Later we will load the SNV data for
    each of these TFBS.
 
    The BED filenames have the pattern
        <jaspar_matrix_id>_<tf_name>.bed

        e.g. MA0651.1_HOXC11.bed
    """

    if bed_file:
        # Directory and BED file specified
        # In case BED file name already includes the path, strip it.
        base_file = os.path.basename(bed_file)
        fullpath = os.path.join(dir, base_file)
        if os.path.isfile(fullpath):
            process_tfbs_and_snv_file(db, fullpath)
        else:
            errstr = "TFBS BED file {} doesn't exist!".format(fullpath)
            logging.error(errstr)
            sys.exit(errstr)
    else:
        # Directory only specified. Process all BED files in directory.
        for file in os.listdir(dir):
            if file.endswith(".bed"):
                process_tfbs_and_snv_file(db, os.path.join(dir, file))


def process_tfbs_and_snv_file(db, bed_file):
    """Process BED file containing the TFBS information and the related
    file containing the SNV information.

    The BED files relate which experiments cover each specific
    TFBS location for the given JASPAR profile.

    These BED files are named:
        <jaspar_matrix_id>_<tf_name>.bed

        e.g. MA0002.2_RUNX1.bed

    The tab-delimited fields contained in these files are:
        chrom, start, end, tf_name, abs_score;rel_score, strand(+/-), exp_id1;exp_id2;...
        e.g. chr1   633930  633941  RUNX1   6.747;83.0  -   GSE64862.runx1.cd34;GSE29180.runx1.jurkat;GSE60130.runx1.amlblast
    """

    bed_basename = os.path.basename(bed_file)
    bed_dirname = os.path.dirname(bed_file)
    bed_rootname = os.path.splitext(bed_basename)[0]

    matrix_id, tf_name = bed_rootname.split('_')[0:2]

    snv_basename = "{}_SNV_impacts.txt".format(bed_rootname)

    snv_file = os.path.join(bed_dirname, matrix_id, snv_basename)

    last_snv_pos = None
    with open(bed_file) as bf, open(snv_file) as sf:
        for bed_line in bf:
            bed_line = bed_line.rstrip()

            (chrom, tfbs_start, tfbs_end, tf_name, score_str, strand,
                exp_str) = bed_line.split('\t')

            chrom = chrom.lstrip('chr')

            abs_score, rel_score = score_str.split(';')

            #
            # Explicitly cast elements to required types for efficient
            # storage in DB.
            #
            tfbs_start = int(tfbs_start) + 1 # convert from BED
                                                # 0-indexed to 1-indexed
                                                # coords
            tfbs_end = int(tfbs_end)
            abs_score = float(abs_score)
            rel_score = float(rel_score) / 100
            
            # experiment list is e.g.:
            #   ENCSR000EGN.jund.k562;GSE51234.jund.gp5d_sirad21
            exp_str = exp_str.upper()
            exp_str_list = exp_str.split(';')

            # The experiment list column sometimes has duplicates,
            # presumably in the case where a TF has two profiles,
            # the mapping was was done via the TF name and not the
            # matrix ID, but this doesn't even seem to be
            # consistent??? Convert to set to remove duplicates.
            exp_str_set = set(exp_str_list)

            exp_id_list = []
            jaspar_tf_name = None
            for exp_str in exp_str_set:
                #exp_id, tf_name, cell_line = exp_str.split('.')

                #curs = db.experiments.find({
                #            'experiment_id' : exp_id,
                #            'cell_line'     : cell_line,
                #            'tf_name'       : tf_name})

                # Retrieve the list of experiments and the JASPAR matrix ID
                # related to this TFBS/SNV. For convenience these are stored
                # in the TFBS/SNV record.
                #for experiment in curs:
                #    exp_id_list.append(experiment['_id'])
                #
                #    for jaspar_tf in experiment['jaspar_tfs']:
                #        if jaspar_tf['matrix_id'] == matrix_id:
                #            jaspar_tf_name = jaspar_tf['name']
                #            break

                # Explicitly check experiment key as there is some bad data
                # in the BED files.
                if exp_str in exp_hash:
                    experiment = exp_hash[exp_str]
                    exp_id_list.append(experiment['_id'])

                    if not jaspar_tf_name:
                        for jaspar_tf in experiment['jaspar_tfs']:
                            if jaspar_tf['matrix_id'] == matrix_id:
                                jaspar_tf_name = jaspar_tf['name']
                                break

            # There shouldn't be duplicates but just to be on the
            # safe side...
            exp_id_set = set(exp_id_list)

            # Now save the basic TFBS information. Later we will load
            # the SNV info from the related SNV impacts file.
            tfbs_snv = {
                'experiment_ids'   : list(exp_id_set),
                'matrix_id'        : matrix_id,
                'jaspar_tf_name'   : jaspar_tf_name,
                'chrom'            : chrom,
                'start'            : tfbs_start,
                'end'              : tfbs_end,
                'abs_score'        : abs_score,
                'rel_score'        : rel_score,
                'strand'           : strand,
                'snvs'             : []
            }

            last_snv_pos = read_and_add_snv_data(sf, tfbs_snv, matrix_id,
                                                 last_snv_pos)

            db.tfbs_snvs.insert(tfbs_snv)


def read_and_add_snv_data(sf, tfbs_snv, matrix_id, last_pos):
    """Read lines from the SNV file while they match the given tfbs_snv
    and add the information to the given tfbs_snv object.
    """

    if last_pos:
        # Last time in the function we read a line for the "next" TFBS/SNV
        # so seek back to it.
        sf.seek(last_pos)

    last_pos = sf.tell()
    snv_line = sf.readline()
    while snv_line:
        snv_line = snv_line.rstrip()

        # NOTE: not a BED file, coords are already 1-based
        (chrom, position, ignore, ref_allele, alt_allele,
        ref_tfbs_start, ref_tfbs_end, ref_tfbs_abs_score,
        ref_tfbs_rel_score, ref_tfbs_strand, snv_matrix_id,
        alt_tfbs_abs_score, alt_tfbs_rel_score, alt_tfbs_strand,
        alt_tfbs_start, alt_tfbs_end, impact_score,
        num_alt_considered, species) = snv_line.split('\t')

        if snv_matrix_id != matrix_id:
            # This shouldn't happen.
            errstr = "Matrix ID '{}' from TFBS/SNV file name {} does not match matrix ID '{}' from TFBS/SNV file line {}".format(
                matrix_id, snv_base_file, snv_matrix_id,
                line_num)

            logging.error(errstr)
            sys.exit(errstr)

        #
        # Explicitly cast elements to required types for efficient
        # storage in DB.
        #
        position = int(position)
        ref_tfbs_start = int(ref_tfbs_start)
        ref_tfbs_end = int(ref_tfbs_end)
        ref_tfbs_abs_score = float(ref_tfbs_abs_score)
        ref_tfbs_rel_score = float(ref_tfbs_rel_score)
        alt_tfbs_abs_score = float(alt_tfbs_abs_score)
        alt_tfbs_rel_score = float(alt_tfbs_rel_score)
        ref_tfbs_strand = '-' if int(ref_tfbs_strand) == -1 else '+'
        alt_tfbs_strand = '-' if int(alt_tfbs_strand) == -1 else '+'
        alt_tfbs_start = int(alt_tfbs_start)
        alt_tfbs_end = int(alt_tfbs_end)
        num_alt_considered = int(num_alt_considered)
        ref_allele = ref_allele.upper()
        alt_allele = alt_allele.upper()

        if tfbs_snv_matches(tfbs_snv, matrix_id, chrom, ref_tfbs_start,
                                ref_tfbs_strand):
            if not tfbs_snv['snvs']:
                # Haven't stored any SNV data for this TFBS/SNV yet.
                snv = {}
                snv['pos'] = position
                snv['ref_allele'] = ref_allele
                snv[alt_allele] = {
                    'start'     : alt_tfbs_start,
                    'end'       : alt_tfbs_end,
                    'strand'    : alt_tfbs_strand,
                    'abs_score' : alt_tfbs_abs_score,
                    'rel_score' : alt_tfbs_rel_score
                }


                if impact_score == 'N/A':
                    # Not sure this is still an issue with MANTA2
                    snv[alt_allele]['impact'] = float(0)
                else:
                    snv[alt_allele]['impact'] = float(impact_score)

                tfbs_snv['snvs'] = [snv]
            else:
                # Add additional SNV data to this TFBS/SNV
                position_exists = False
                for snv in tfbs_snv['snvs']:
                    if snv['pos'] == position:
                        # An SNV already exists at this  position.
                        # Add new alt. allele data to it.
                        position_exists = True

                        if alt_allele in snv:
                            #
                            # XXX
                            # This should not happen.
                            # We do have duplicates. Not sure why.
                            # XXX
                            #
                            errstr = "Duplicate alt. allele {}, position {}, ref. allele {} in file {}, line {}:\n{}\nexisting: {}".format(alt_allele, position, ref_allele, snv_file, line_num, snv_line, snv)
                            logging.error(errstr)
                            #sys.exit(errstr)
                        else:
                            snv[alt_allele] = {
                                'start'     : alt_tfbs_start,
                                'end'       : alt_tfbs_end,
                                'strand'    : alt_tfbs_strand,
                                'abs_score' : alt_tfbs_abs_score,
                                'rel_score' : alt_tfbs_rel_score
                            }

                            if impact_score == 'N/A':
                                snv[alt_allele]['impact'] = float(0)
                            else:
                                snv[alt_allele]['impact'] = float(impact_score)
                        break

                if not position_exists:
                    # There was not existing SNV for this position.
                    # Add SNV for this new position within the binding site.
                    snv = {
                        'pos'           : position,
                        'ref_allele'    : ref_allele,
                        alt_allele : {
                            'start'     : alt_tfbs_start,
                            'end'       : alt_tfbs_end,
                            'strand'    : alt_tfbs_strand,
                            'abs_score' : alt_tfbs_abs_score,
                            'rel_score' : alt_tfbs_rel_score,
                        }
                    }

                    if impact_score == 'N/A':
                        snv[alt_allele]['impact'] = float(0)
                    else:
                        snv[alt_allele]['impact'] = float(impact_score)

                    tfbs_snv['snvs'].append(snv)

            last_pos = sf.tell()
            snv_line = sf.readline()
        else:
            # We hit a line for a different TFBS/SNV (or end of file).
            # Record last position and return.
            return last_pos
    
    return None


#def load_snvs(db, dir, bed_file=None):
#    """Load the TFBS/SNV data. This stores the SNV information which consists
#    of all the computational mutations at each position and the changes these
#    result in, regarding the matrix score, binding site shift etc. The main
#    TFBS information has already been stored in the table in the load_tfbs
#    routine, so we are adding the SNV information to each existing record.
#    """
#
#    if bed_file:
#        # Directory and file specified
#        # In case file name already includes the path, strip it.
#        base_file = os.path.basename(bed_file)
#        root_file = os.path.splitext(base_file)[0] # without ext.
#
#        matrix_id, tf_name = root_file.split('_')[0:2]
#
#        #MA0685.1_SP4_SNV_impacts.txt
#
#        impacts_file = "{}_{}_SNV_impacts.txt".format(matrix_id, tf_name)
#
#        fullpath = os.path.join(dir, matrix_id, impacts_file)
#        if os.path.isfile(fullpath):
#            process_snv_file(db, fullpath)
#        else:
#            errstr = "SNV impacts file {} doesn't exist!".format(fullpath)
#            logging.error(errstr)
#            sys.exit(errstr)
#    else:
#        # Directory only specified. Process all SNV subdirectories/files
#        # in this directory.
#        for file in os.listdir(dir):
#            #
#            # The base directory contains subdirectories for each JASPAR motif
#            # which use the matrix ID as the sub dir name, as well as BED files
#            # containing the experiment to TFBS mapping files.
#            #
#            fullpath = os.path.join(dir, file)
#
#            if os.path.isdir(fullpath):
#                for f in os.listdir(fullpath):
#                    if fnmatch.fnmatch(f, "{}_*_impacts.txt".format(file)):
#                        process_snv_file(db, os.path.join(fullpath, f))
#
#
#def process_snv_file(db, snv_file):
#    snv_base_file = os.path.basename(snv_file) # without path
#    snv_root_file = os.path.splitext(snv_base_file)[0] # without ext.
#    matrix_id, tf_name = snv_root_file.split('_')[0:2]
#
#    tfbs_snv = None
#    line_num = 0
#
#    with open(snv_file, 'r') as f:
#        for snv_line in f:
#            snv_line = snv_line.rstrip()
#            line_num += 1
#
#            # NOTE: not a BED file, coords are already 1-based
#            (chrom, position, ignore, ref_allele, alt_allele,
#            ref_tfbs_start, ref_tfbs_end, ref_tfbs_abs_score,
#            ref_tfbs_rel_score, ref_tfbs_strand, snv_matrix_id,
#            alt_tfbs_abs_score, alt_tfbs_rel_score, alt_tfbs_strand,
#            alt_tfbs_start, alt_tfbs_end, impact_score,
#            num_alt_considered, species) = snv_line.split('\t')
#
#            if snv_matrix_id != matrix_id:
#                errstr = "Matrix ID '{}' from TFBS/SNV file name {} does not match matrix ID '{}' from TFBS/SNV file line {}".format(
#                    matrix_id, snv_base_file, snv_matrix_id,
#                    line_num)
#
#                logging.error(errstr)
#                sys.exit(errstr)
#
#            #
#            # Explicitly cast elements to required types for efficient
#            # storage in DB.
#            #
#            position = int(position)
#            ref_tfbs_start = int(ref_tfbs_start)
#            ref_tfbs_end = int(ref_tfbs_end)
#            ref_tfbs_abs_score = float(ref_tfbs_abs_score)
#            ref_tfbs_rel_score = float(ref_tfbs_rel_score)
#            alt_tfbs_abs_score = float(alt_tfbs_abs_score)
#            alt_tfbs_rel_score = float(alt_tfbs_rel_score)
#            ref_tfbs_strand = '-' if int(ref_tfbs_strand) == -1 else '+'
#            alt_tfbs_strand = '-' if int(alt_tfbs_strand) == -1 else '+'
#            alt_tfbs_start = int(alt_tfbs_start)
#            alt_tfbs_end = int(alt_tfbs_end)
#            num_alt_considered = int(num_alt_considered)
#            ref_allele = ref_allele.upper()
#            alt_allele = alt_allele.upper()
#
#            if tfbs_snv:
#                # Already started a TFBS/SNV
#                if tfbs_snv_matches(tfbs_snv, matrix_id, chrom, ref_tfbs_start,
#                                    ref_tfbs_strand):
#                    # Same TFBS. Add additional SNV impact data to existing
#                    # site.
#
#                    position_exists = False
#                    for snv in tfbs_snv['snvs']:
#                        if snv['pos'] == position:
#                            # Same position. Add new alt. allele data.
#                            position_exists = True
#
#                            if alt_allele in snv:
#                                #
#                                # XXX
#                                # This should not happen.
#                                # We do have duplicates. Not sure why.
#                                # XXX
#                                #
#                                errstr = "Duplicate alt. allele {}, position {}, ref. allele {} in file {}, line {}:\n{}\nexisting: {}".format(alt_allele, position, ref_allele, snv_file, line_num, snv_line, snv)
#                                logging.error(errstr)
#                                #sys.exit(errstr)
#                            else:
#
#                                snv[alt_allele] = {
#                                    'start'     : alt_tfbs_start,
#                                    'end'       : alt_tfbs_end,
#                                    'strand'    : alt_tfbs_strand,
#                                    'abs_score' : alt_tfbs_abs_score,
#                                    'rel_score' : alt_tfbs_rel_score
#                                }
#
#                                if impact_score != 'N/A':
#                                    snv[alt_allele]['impact'] = float(impact_score)
#
#                            break
#
#                    if not position_exists:
#                        # New position within the binding site.
#                        snv = {
#                            'pos'           : position,
#                            'ref_allele'    : ref_allele,
#                            alt_allele : {
#                                'start'     : alt_tfbs_start,
#                                'end'       : alt_tfbs_end,
#                                'strand'    : alt_tfbs_strand,
#                                'abs_score' : alt_tfbs_abs_score,
#                                'rel_score' : alt_tfbs_rel_score,
#                            }
#                        }
#
#                        if impact_score != 'N/A':
#                            snv[alt_allele]['impact'] = float(impact_score)
#
#                        tfbs_snv['snvs'].append(snv)
#                else:
#                    # This is a different TFBS/SNV.
#                    # Save current tfbs_snv and start new one.
#                    db.tfbs_snvs.update(
#                        {
#                            'matrix_id' : tfbs_snv['matrix_id'],
#                            'chrom'     : tfbs_snv['chrom'],
#                            'start'     : tfbs_snv['start'],
#                            'strand'    : tfbs_snv['strand']
#                        },
#                        tfbs_snv
#                    )
#
#                    # Find the existing partial tfbs_snv record for the
#                    # current line being processed that was previously saved
#                    # in the DB by the load_tfbs routine.
#                    curs = db.tfbs_snvs.find({
#                        'matrix_id'     : matrix_id,
#                        'chrom'         : chrom,
#                        'start'         : ref_tfbs_start,
#                        'strand'        : ref_tfbs_strand,
#                    })
#
#                    if not curs or curs.count() == 0:
#                        errstr = "No existing tfbs_snv record found for {}|{}|{}|{} in file {} at line {}\n".format(matrix_id, chrom, ref_tfbs_start, ref_tfbs_end, snv_file, line_num)
#                        logging.error(errstr)
#                        sys.exit(errstr)
#                    elif curs.count() > 1:
#                        errstr = "More than one existing tfbs_snv record found for {}|{}|{}|{} in file {} at line {}\n".format(matrix_id, chrom, ref_tfbs_start, ref_tfbs_end, snv_file, line_num)
#                        logging.error(errstr)
#                        sys.exit(errstr)
#
#                    tfbs_snv = curs.next()
#                    tfbs_snv['alt_considered'] = num_alt_considered
#                    tfbs_snv['snvs'] = []
#
#                    snv = {
#                        'pos'           : position,
#                        'ref_allele'    : ref_allele,
#                        alt_allele : {
#                            'start'     : alt_tfbs_start,
#                            'end'       : alt_tfbs_end,
#                            'strand'    : alt_tfbs_strand,
#                            'abs_score' : alt_tfbs_abs_score,
#                            'rel_score' : alt_tfbs_rel_score
#                        }
#                    }
#
#                    if impact_score != 'N/A':
#                        snv[alt_allele]['impact'] = float(impact_score)
#
#                    tfbs_snv['snvs'].append(snv)
#            else:
#                # No currrent tfbs_snv. Start a new one. We should be at the
#                # first line of the file. Fetch the existing partial tfbs_snv
#                # record that has already been saved in the DB via the
#                # load_tfbs routine.
#                curs = db.tfbs_snvs.find({
#                    'matrix_id'     : matrix_id,
#                    'chrom'         : chrom,
#                    'start'         : ref_tfbs_start,
#                    'strand'        : ref_tfbs_strand,
#                })
#
#                # Add info
#                if not curs or curs.count() == 0:
#                    errstr = "No existing tfbs_snv record found for {}|{}|{}|{} in file {} at line {}\n".format(matrix_id, chrom, ref_tfbs_start, ref_tfbs_end, snv_file, line_num)
#                    logging.error(errstr)
#                    sys.exit(errstr)
#                elif curs.count() > 1:
#                    errstr = "More than one existing tfbs_snv record found for {}|{}|{}|{} in file {} at line {}\n".format(matrix_id, chrom, ref_tfbs_start, ref_tfbs_end, snv_file, line_num)
#                    logging.error(errstr)
#                    sys.exit(errstr)
#
#                tfbs_snv = curs.next()
#                tfbs_snv['alt_considered'] = num_alt_considered
#                tfbs_snv['snvs'] = []
#
#                snv = {
#                    'pos'           : position,
#                    'ref_allele'    : ref_allele,
#                    alt_allele : {
#                        'start'     : alt_tfbs_start,
#                        'end'       : alt_tfbs_end,
#                        'strand'    : alt_tfbs_strand,
#                        'abs_score' : alt_tfbs_abs_score,
#                        'rel_score' : alt_tfbs_rel_score
#                    }
#                }
#
#                if impact_score != 'N/A':
#                    snv[alt_allele]['impact'] = float(impact_score)
#
#                tfbs_snv['snvs'].append(snv)
#
#    # Save final TFBS_SNV
#    db.tfbs_snvs.update(
#        {
#            'matrix_id' : tfbs_snv['matrix_id'],
#            'chrom'     : tfbs_snv['chrom'],
#            'start'     : tfbs_snv['start'],
#            'strand'    : tfbs_snv['strand']
#        }, tfbs_snv)
        

def tfbs_snv_matches(tfbs_snv, matrix_id, chrom, start, strand):
    return (chrom == tfbs_snv['chrom'] and start == tfbs_snv['start']
            and strand == tfbs_snv['strand'])


def fetch_experiments_into_hash(db):
    """Read all records in experiments collection and store in a hash
    """

    exp_hash = {}

    for exp in db.experiments.find():
        # Use the same format as the experiment strings in the BED files
        # except here everything is uppercase.
        exp_key = "{}.{}.{}".format(exp['experiment_id'],
                                    exp['tf_name'],
                                    exp['cell_line'])

        if exp_key in exp_hash:
            sys.exit("Duplicate experiment for hash_key {}".format(exp_key))

        exp_hash[exp_key] = exp

    return exp_hash


###############################################################################
#                               MAIN
###############################################################################
if __name__ == '__main__':
    """Load the MANTA database.

    Usage: load_database.py -d mant_data_dir [-c]
                [-u db_user] -p db_pass [-xo] [-f bed_file]

    Where:
        -d DIR  - Specifies the top level directory containing the master
                  CSV file, BED files and subdirecotories with the TFBS/SNV
                  data.
        -c      - If specified, the database is completely cleaned of all
                  existing documents.
        -u user - MANTA2 db username
        -p pswd - MANTA2 db password (must have write access)
        -xo     - Indicates loading the experiments only.
        -f      - Input BED file. Process only this BED file coorresponding
                  a specific TF (see below).

    If the script is run in a "one-shot" manner from scratch, give the
    -d and -c (and optionally -u and -p) parameters to build the entire
    MANTA2 database in a single process.

    If running on the cluster, first run with the -d, -c and -xo (and
    -u and -p if needed) parameters to first build the experiments
    collection.

    Then run individual jobs for each TF on the cluster nodes by using the
    -d and -f (and -u and -p if needed) parameters for each individual
    BED file containing the TFBS information.
    """

    parser = argparse.ArgumentParser(
        description='Parse the various files associated with the TFBS/SNVs and load the information into the MANTA2 database'
    )

    parser.add_argument(
        '-d', '--dir', nargs='?', required=True, help='Top level directory containing the master CSV file, BED files with the TFBS to experment relation BED files and the subdirectories containing the text files with all the TFBS/SNVs.'
    )

    parser.add_argument(
        '-u', '--user', nargs='?', default=MANTA2_DB_USER, const=MANTA2_DB_USER, help='MANTA2 DB user name'
    )

    parser.add_argument(
        '-p', '--password', nargs='?', required=True, help='MANTA2 DB password'
    )

    parser.add_argument(
        '-xo', '--experiments-only', action='store_true', help='If specified, load only the experiments collection'
    )

    parser.add_argument(
        '-f', '--bed-file', nargs='?', help='If specified, only load TFBS/SNV based on this input BED file. This assumes the experiments collection has already been loaded into the DB'
    )

    parser.add_argument(
        '-c', '--clean_db', action='store_true', help='If specified, clean the database before update'
    )

    parser.add_argument(
        '-l', '--log', nargs='?', const=LOGFILE, default=LOGFILE, help='Name of file to which logging messages are written'
    )

    args = parser.parse_args()

    dir=args.dir
    manta2_db_user = args.user
    manta2_db_pass = args.password
    logfile = args.log
    clean_db = args.clean_db
    exp_only = args.experiments_only
    bed_file = args.bed_file

    if exp_only and bed_file:
        sys.exit("The --experiments-only and --bed-file are exlusive options as running with a single bed file indicates that you are running a process to load TFBS/SNV data for a single TF, and assumes experiments have already been loaded in a previous run with the --experiments-only flag.")

    do_load_exp = False
    if exp_only or not bed_file:
        do_load_exp = True

    do_load_tfbs_snv = False
    if not exp_only or bed_file:
        do_load_tfbs_snv = True

    if do_load_exp and not clean_db:
        sys.exit("Provided options indicate populating the experiments collection will take place but the clean-db option was not specified. Generally, loading experiments indicates the database is being built from scratch and any existing data should be cleaned by specifying the clean-db option.")

    logging.basicConfig(filename=logfile, filemode='w', level=logging.INFO)

    start_datetime = datetime.datetime.now()
    logging.info("MANTA2 DB build started on {0}\n".format(start_datetime))
    logging.info("MANTA2 top level data directory: {0}\n".format(dir))

    uri = "mongodb://{0}:{1}@{2}/{3}".format(
        manta2_db_user, manta2_db_pass, MANTA2_DB_HOST, MANTA2_DB_NAME
    )
    client = MongoClient(uri)
    db = client.manta2

    #
    # If -c option specified, clean database (remove all documents from all
    # collections) before insert/update of new data.
    #
    if clean_db:
        clean_database(db)

    if do_load_exp:
        load_experiments(db, dir)

    exp_hash = fetch_experiments_into_hash(db)

    if do_load_tfbs_snv:
        #load_tfbs(db, dir, bed_file=bed_file)
        #load_snvs(db, dir, bed_file=bed_file)
        load_tfbs_snv(db, dir, bed_file=bed_file)

    client.close()

    end_datetime = datetime.datetime.now()
    elapsed_time = end_datetime - start_datetime
    logging.info("MANTA DB build completed successfully on {0}\n".format(end_datetime))
    logging.info("Elapsed time: {0}\n".format(elapsed_time))
