#!/usr/bin/env python

import os;

import sys;

import string;

from collections import defaultdict

import argparse;

from argparse import RawTextHelpFormatter

from DeepMod_scripts.myCom import *

# three modules in DeepMod

parser = argparse.ArgumentParser(description="Detect nucleotide modification from nanopore signals data.", epilog="For example, \n \

\tpython %(prog)s train: Training a modification classifier.\n \

\tpython %(prog)s detect: Detect modification by integrating all long reads. \n \

\tpython %(prog)s getfeatures: Get features for training a model.  \n \

", formatter_class=RawTextHelpFormatter);

#

# Return error message when a value<1

# Return an empty string otherwise

#

def non_negative(i, mstr):

   if i<1: return (("\n\tError %d could not be negative(%d)" % (mstr, i)))

   else: return ''

#

# Print all parameters in stdout

#

def printParameters(moptions):

   mpkeys = moptions.keys(); #mpkeys.sort()

   sorted(mpkeys)

   print('%30s: %s' % ('Current directory', os.getcwd()))

   for mpk in mpkeys:

      print ('%30s: %s' % (mpk, str(moptions[mpk])))

   sys.stdout.flush()

#

# Got common argument provided by users or default values.

#

#

def mCommonParam(margs):

   ErrorMessage = ""

   moptions = defaultdict()

   # how to output running message: need more control now.

   moptions['outLevel'] = margs.outLevel

   # the input working base

   moptions["wrkBase"] = margs.wrkBase

   if moptions["wrkBase"]==None:

      ErrorMessage = ErrorMessage + ("\n\tThe input folder is None.")

   # An unique ID for output

   # Usefull for run the program in parallel

   moptions["FileID"] = margs.FileID

   # output folder;

   # make it if the output folder does not exist

   moptions['outFolder'] = margs.outFolder

   moptions['outFolder'] = format_last_letter_of_folder(moptions['outFolder'])

   if moptions['outFolder']==None or (not os.path.isdir(moptions['outFolder'])):

      try:

         os.system('mkdir -p '+moptions['outFolder']);

      except:

         ErrorMessage = ErrorMessage + ("\n\tThe output folder (%s) does not exist and cannot be created." % moptions['outFolder'])

   # check all data in a recurive way

   moptions['recursive'] = margs.recursive

   # the number of threads used and the number of files handled by each thread.

   moptions['files_per_thread'] = margs.files_per_thread

   if moptions['files_per_thread']<2: moptions['files_per_thread'] = 2

   # the number of threads used

   moptions['threads'] = margs.threads

   if moptions['threads']<1: moptions['threads'] = 1

   # windowsize: default=21

   moptions['windowsize'] = margs.windowsize

   ErrorMessage = ErrorMessage + non_negative(moptions['windowsize'], 'windowsize')

   if moptions['windowsize']<1: moptions['windowsize'] = 1

   # aligners: bwa-mem or minimap2

   moptions['alignStr'] = margs.alignStr;

   moptions['SignalGroup'] = margs.SignalGroup;

   moptions['move'] = margs.move

   return [moptions, ErrorMessage]

#

# detect modification for bases of interests

# input is a list of fast5 files, a reference genome and a well-trained model.

#

def mDetect(margs):

   # get common parameters

   moptions, ErrorMessage = mCommonParam(margs)

   # path for basecall information in fast5 files

   moptions['basecall_1d'] = margs.basecall_1d

   moptions['basecall_2strand'] = margs.basecall_2strand

   # Whether consider those chromosome which contain -_:/

   # default: yes;

   moptions['ConUnk'] = margs.ConUnk

   # output layer information for deep learning

   moptions['outputlayer'] = margs.outputlayer

   # base of interest

   moptions['Base'] = margs.Base

   # whether take cluster effect of methylation into consideration

   moptions['mod_cluster'] = margs.mod_cluster

   # base of interest

   if moptions['Base'] in ["", None]:

      ErrorMessage = ErrorMessage + ("\n\t Please provide a base of interest.")

   # predict medification for bases of interest in long reads first

   # only summarize them for each genomic position of interest .

   moptions['predDet'] = margs.predDet

   if moptions['predDet']:

      # path to reference genome

      moptions['Ref'] = margs.Ref

      if moptions['Ref']==None or (not os.path.isfile(moptions['Ref'])):

         ErrorMessage = ErrorMessage + ("\n\t reference file does not exist (%s)" % moptions['Ref'])

      # the number of feature for each event

      moptions['fnum'] = margs.fnum

      ErrorMessage = ErrorMessage + non_negative(moptions['fnum'], 'fnum')

      # the number of hidden nodes

      moptions['hidden'] = margs.hidden

      ErrorMessage = ErrorMessage + non_negative(moptions['hidden'], 'hidden')

      # the well-trained model

      moptions['modfile'] = margs.modfile

      if moptions['modfile']==None:

         print("No mod file is provided. The default one is used")

         moptions['modfile'] = ('train_deepmod/rnn_P90wd%d_f53/mod_train_P90wd%d_f53' % (moptions['windowsize'], moptions['windowsize']))

         if (not os.path.isfile(moptions['modfile']+'.meta')):

            moptions['modfile'] = ('{}/lib/python{}.{}/site-packages/DeepMod/train_deepmod/rnn_P90wd{}_f53/mod_train_P90wd{}_f53'.format(sys.prefix,sys.version_info.major,sys.version_info.minor, moptions['windowsize'], moptions['windowsize']))

      if (not os.path.isfile(moptions['modfile']+'.meta')):

         ErrorMessage = ErrorMessage + ("\n\tThe meta file (%s) does not exist" % (moptions['modfile']+'.meta' if not moptions['modfile']==None else ""))

   else:

      # already done the prediction process?

      # Yes: summarize the results only

      moptions['predpath'] = margs.predpath

      if moptions['predpath']==None or (not os.path.isdir(moptions['predpath'])):

         ErrorMessage = ErrorMessage + ("\n\tThe predpath does not exist")

   # specify region of interest

   # not consider bases outside regions in a reference genome

   # None: all bases of interest

   moptions['region'] = [ ]

   if margs.region == None or len(margs.region)==0:

      moptions['region'].append([None, None, None])

   else:

      mregionlist = margs.region.split(';')

      for mr in mregionlist:

         mr_sp = mr.split(':')

         moptions['region'].append([mr_sp[0], int(mr_sp[1]) if len(mr_sp)>1 else None, int(mr_sp[2]) if len(mr_sp)>2 else None ])

   print("\nNanopore sequencing data analysis is resourece-intensive and time consuming. ")

   print("Some potential strong recommendations are below:")

   print("\tIf your reference genome is large as human genome and your Nanopore data is huge,")

   print("\tIt would be faster to run this program parallelly to speed up.")

   print("\tYou might run different input folders of your fast5 files and ")

   print("\tgive different output names (--FileID) or folders (--outFolder)")

   print("\tA good way for this is to run different chromosome individually.\n")

   # print help information if any necessary options are not provided.

   printParameters(moptions)

   if not ErrorMessage=="":

      ErrorMessage = "Please provide correct parameters" + ErrorMessage

      print(ErrorMessage)

      parser.print_help();

      parser.parse_args(['detect', '--help']);

      sys.exit(1)

   from DeepMod_scripts import myDetect

   myDetect.mDetect_manager(moptions)

#

# Train a model

# Need to get features first.

#

def mTrain(margs):

   from DeepMod_scripts import myMultiBiRNN

   # gent common options

   moptions, ErrorMessage = mCommonParam(margs)

   # network setting: the number of features and the number of hidden nodes

   moptions['fnum'] = margs.fnum

   ErrorMessage = ErrorMessage + non_negative(moptions['fnum'], 'fnum')

   moptions['hidden'] = margs.hidden

   ErrorMessage = ErrorMessage + non_negative(moptions['hidden'], 'hidden')

   # the output function of the deep learning model

   moptions['outputlayer'] = margs.outputlayer

   # whether using different class weights

   moptions['unbalanced'] = margs.unbalanced

   # re-load trained model and continue to train

   moptions['modfile'] = margs.modfile

   if moptions['modfile']==None: pass;

   elif (not os.path.isfile(moptions['modfile']+'.meta')):

      ErrorMessage = ErrorMessage + ("\n\tThe meta file (%s) does not exist" % (moptions['modfile']+'.meta' if not moptions['modfile']==None else ""))

   # read-based or region based independent training

   # E: region-based

   # P: read-based.

   if not margs.test==None:

      moptions['test'] = margs.test.split(',')

      if moptions['test'][0] == 'E': moptions['test'][0] = '-'

      elif moptions['test'][0] == 'P': moptions['test'][0] = '0'

      else:

         ErrorMessage = ErrorMessage + "Unknown option for test: the first character must be E or P "+margs.test

      if moptions['test'][0] in ['-']:

         moptions['test'][1] = int(moptions['test'][1]) * (10**6)

         moptions['test'][2] = int(moptions['test'][2]) * (10**6)

      else: moptions['test'][1] = int(moptions['test'][1])/100.0

   else: moptions['test'] = ['N', '100']

   # print help document if necessary options are not provided.

   print("Train")

   printParameters(moptions)

   if not ErrorMessage=="":

      ErrorMessage = "Please provide correct parameters" + ErrorMessage

      print(ErrorMessage)

      parser.print_help();

      parser.parse_args(['train', '--help']);

      sys.exit(2)

   myMultiBiRNN.mMult_RNN_LSTM_train(moptions)

#

# get features for training

#

#

def mGetFeatures(margs):

   from DeepMod_scripts import myGetFeatureBasedPos

   # get common options

   moptions, ErrorMessage = mCommonParam(margs)

   # motif-based data: positive or negative control data

   moptions['posneg'] = margs.posneg

   # the number of features: 7-description or 57-description

   moptions['fnum'] = margs.fnum

   ErrorMessage = ErrorMessage + non_negative(moptions['fnum'], 'fnum')

   # size of each bacth to store features

   moptions['size_per_batch'] = margs.size_per_batch

   if moptions['size_per_batch'] < 0.001: moptions['size_per_batch'] = 0.001

   # path to basecall inform in fast5 files

   moptions['basecall_1d'] = margs.basecall_1d

   moptions['basecall_2strand'] = margs.basecall_2strand

   # regions of interest

   moptions['region'] = [None, None, None]

   if not (margs.region==None or margs.region.strip()==''):

      rsp = margs.region.split(':')

      for rv_ind in range(len(rsp)):

         rsp[rv_ind] = rsp[rv_ind].strip();

         if not rsp[rv_ind]=='':

            moptions['region'][rv_ind] = rsp[rv_ind]

   # referene genome

   moptions['Ref'] = margs.Ref

   if moptions['Ref']==None or (not os.path.isfile(moptions['Ref'])):

      ErrorMessage = ErrorMessage + ("\n\t reference file does not exist (%s)" % moptions['Ref'])

   # get motif-based modification

   # or specify by --fulmod/--anymod/--nomod

   moptions['motifORPos'] = margs.motifORPos

   if margs.motifORPos==1:

      moptions['motif'] = [margs.motif.upper(), margs.ModinMotif]

   elif margs.motifORPos==2:

      moptions['fulmod'] = margs.fulmod

      if moptions['fulmod']==None: # completely modificated positions

         ErrorMessage = ErrorMessage + ("\t There is no parameter for --fulmod.")

      moptions['anymod'] = margs.anymod

      if moptions['anymod'] == None: # patially modificated positions

         ErrorMessage = ErrorMessage + ("\t There is no parameter for --anymod.")

      moptions['nomod'] = margs.nomod

      if moptions['nomod'] == None: # completely unmodified posisionts

         ErrorMessage = ErrorMessage + ("\t There is no parameter for --nomod.")

   else:

      ErrorMessage = ErrorMessage + ("\tmotifORPos value (%d) is not supported." % margs.motifORPos)

   # print help document if any required options are not provided.

   printParameters(moptions)

   if not ErrorMessage=="":

      ErrorMessage = "Please provide correct parameters" + ErrorMessage

      print(ErrorMessage)

      parser.print_help();

      parser.parse_args(['getfeatures', '--help']);

      sys.exit(1)

   myGetFeatureBasedPos.getFeature_manager(moptions)

#####################################################################################

subparsers = parser.add_subparsers()

parent_parser = argparse.ArgumentParser(add_help=False)

# add common options

com_group_for_comparison = parent_parser.add_argument_group('Common options.')

com_group_for_comparison.add_argument("--outLevel", type=int, choices=[OUTPUT_DEBUG, OUTPUT_INFO, OUTPUT_WARNING, OUTPUT_ERROR], default=OUTPUT_WARNING, help=("The level for output: %d for DEBUG, %d for INFO, %d for WARNING, %d for ERROR. Default: %d" % (OUTPUT_DEBUG, OUTPUT_INFO, OUTPUT_WARNING, OUTPUT_ERROR, OUTPUT_WARNING)))

com_group_for_comparison.add_argument("--wrkBase", help="The base folder for FAST5 files.")

com_group_for_comparison.add_argument("--FileID", default="mod", help="The unique string for intermediate files and final output files. Default: 'mod'")

com_group_for_comparison.add_argument("--outFolder", default='./mod_output', help="The default folder for outputing the results. Default: ./mod_output")

com_group_for_comparison.add_argument("--recursive", type=int, default=1, choices=[0,1], help="Recurise to find fast5 files. Default:1")

com_group_for_comparison.add_argument("--threads", type=int, default=4, help="The number of threads used (not for train). Default:4")

com_group_for_comparison.add_argument("--files_per_thread", type=int, default=1000, help="The number of fast5 files for each thread (not for train). Default:500")

com_group_for_comparison.add_argument("--windowsize", type=int, default=21, help="The window size to extract features. Default: 21")

com_group_for_comparison.add_argument("--alignStr", type=str, default='minimap2', choices=["bwa","minimap2"], help="Alignment tools (bwa or minimap2 is supported). Default: minimap2")

com_group_for_comparison.add_argument("--SignalGroup", type=str, default='simple', choices=["simple","rundif"], help="How to associate signals to each called bases. Default: simple")

com_group_for_comparison.add_argument("--move", default=False, action="store_true", help="Whether the basecalled data use move tables instead of event tables. Default: False")

# add detection options

parser_detect = subparsers.add_parser('detect', parents=[parent_parser], help="Detect modifications at a genomic scale", description="Detect modifications by integrating all long reads for a genome", epilog="For example, \n \

python %(prog)s --wrkBase ctrl_oligo_SpeI_cut --FileID mod_det --outFolder ./mod_output/detect3 \n \

", formatter_class=RawTextHelpFormatter)

parser_detect.add_argument("--Ref", help="The reference sequence")

parser_detect.add_argument("--predDet", type=int, default=1, choices=[0,1], help="pred first and then detect (1) or only detect (0). Default: 1")

parser_detect.add_argument("--predpath", default=None, help="The file path of predictions for each fast5 file. The file pattern is *_*.detail. Default: './mod_output/pred2/'")

parser_detect.add_argument("--modfile", type=str, default=None, help="The path to load training model. Default: 'mod_output/'")

parser_detect.add_argument("--fnum", type=int, default=7, help="The number of features. Default: 7")

parser_detect.add_argument("--hidden", type=int, default=100, help="The number of hidden node. Default: 100")

parser_detect.add_argument("--basecall_1d", default="Basecall_1D_000", help="Path for basecall_1d. Default: Basecall_1D_000")

parser_detect.add_argument("--basecall_2strand", default="BaseCalled_template", help="Path for basecall_2strand. Default: BaseCalled_template")

parser_detect.add_argument("--region", default=None, help="The region of interest: for example, chr:1:100000;chr2:10000");

parser_detect.add_argument("--ConUnk", default=True, choices=[False, True], help="Whether contain unknown chromosome");

parser_detect.add_argument("--outputlayer", default="", choices=["", "sigmoid"], help="how to put activation function for output layer")

parser_detect.add_argument("--Base", type=str, default='C', choices=['A', 'C', 'G', 'T'], help="Interest of bases");

parser_detect.add_argument("--mod_cluster", default=0, choices=[0,1], help="1: CpG cluster effect; 0: not");

parser_detect.set_defaults(func=mDetect)

# add training options

parser_training = subparsers.add_parser('train', parents=[parent_parser], help="Training a modification classifier", description="Training a modification classifier", epilog="For example, \n \

python %(prog)s --wrkBase umr --wrkBase2 sss --FileID mod_train --outFolder ./mod_output/train1 \n \

", formatter_class=RawTextHelpFormatter)

parser_training.add_argument("--wrkBase2", help="The base folder for long reads without any modifications.")

parser_training.add_argument("--fnum", type=int, default=7, help="The number of features. Default: 7")

parser_training.add_argument("--hidden", type=int, default=100, help="The number of hidden node. Default: 100")

parser_training.add_argument("--modfile", type=str, default=None, help="The path to load training model. Default: 'mod_output/'")

parser_training.add_argument("--test", help="The number of E Coli genomic position for testing. Default: 'E,1,2'")

parser_training.add_argument("--outputlayer", default="", choices=["", "sigmoid"], help="how to put activation function for output layer")

parser_training.add_argument("--unbalanced", type=int, default=0, choices=[1, 0, None], help="Whether data is unbalanced");

parser_training.set_defaults(func=mTrain)

# add get-feature options

parser_getfeatures = subparsers.add_parser('getfeatures', parents=[parent_parser], help="Get features for all fast5 files", description="Get features for all fast5 files", epilog="For example, \n \

python %(prog)s --wrkBase umr/160617_ecolilowinput_UMR9/called/pass --threads 48 --recursive 0 --posneg 0 --outFolder umr  \n \

python %(prog)s --wrkBase sss/160617_ecolilowinput_sssiR9/called/pass --threads 48 --recursive 0 --posneg 1 --outFolder sss \n \

", formatter_class=RawTextHelpFormatter)

parser_getfeatures.add_argument("--posneg", type=int, default=0, choices=[0,1], help="The positive(1) or negative(0) class. Default: 0")

parser_getfeatures.add_argument("--size_per_batch", type=int, default=1, help="The size (unit: 10^7=10M) of a feature file. Default: 1")

parser_getfeatures.add_argument("--fnum", type=int, default=7, help="The number of features. Default: 7")

parser_getfeatures.add_argument("--region", type=str, help="The region of interest. Set to None or empty for all. Format is chr:start_pos:end_pos")

parser_getfeatures.add_argument("--basecall_1d", default="Basecall_1D_000", help="Path for basecall_1d. Default: Basecall_1D_000")

parser_getfeatures.add_argument("--basecall_2strand", default="BaseCalled_template", help="Path for basecall_2strand. Default: BaseCalled_template")

parser_getfeatures.add_argument("--motifORPos", type=int, default=1, help="Use Motif (1) or pos (2) for modified bases. Default: 1")

parser_getfeatures.add_argument("--motif", default='CG', type=str, help="The motif of interest")

parser_getfeatures.add_argument("--ModinMotif", default=0, type=int, help="The position of modified base in the motif of interest")

parser_getfeatures.add_argument("--Ref", help="The reference sequence")

parser_getfeatures.add_argument("--fulmod", type=str, help="The file pattern for full modification: bisultfiteseq/chr20_C*_0.95.txt")

parser_getfeatures.add_argument("--anymod", type=str, help="The file pattern for any modification: bisultfiteseq/chr20_any_0.95.txt")

parser_getfeatures.add_argument("--nomod", type=str, help="The file pattern for any modification: bisultfiteseq/chr20_no1_0.95.txt")

parser_getfeatures.set_defaults(func=mGetFeatures)

# no provided argument

# print help document

if len(sys.argv)<2:

   parser.print_help();

else:

   args = parser.parse_args()

   args.func(args);