#!/usr/bin/env python
# encoding: utf-8


import sys
import os
from argparse import ArgumentParser

parser = ArgumentParser('script to create a DeepJetCore subpackage')

parser.add_argument("subpackage_name", help="name of the subpackage")
parser.add_argument("subpackage_parent_dir", help="parent directory of the subpackage (must be same as DeepJetCore)")
parser.add_argument("--nodata", help="do NOT create example data", default=False, action="store_true")

args=parser.parse_args()

deepjetcore = os.getenv('DEEPJETCORE')

subpackage_dir=args.subpackage_parent_dir+'/'+args.subpackage_name

### templates ####

environment_file='''
#! /bin/bash
THISDIR=`pwd`
cd {deepjetcore}
source env.sh
cd $THISDIR
export {subpackage}=$( cd "$( dirname "${BASH_SOURCE}" )" && pwd -P)
export DEEPJETCORE_SUBPACKAGE=${subpackage}

cd ${subpackage}
export PYTHONPATH=${subpackage}/modules:$PYTHONPATH
export PYTHONPATH=${subpackage}/modules/datastructures:$PYTHONPATH
'''.format(deepjetcore=deepjetcore, 
           subpackage=args.subpackage_name.upper(),
           subpackage_dir=os.path.abspath(subpackage_dir),
           BASH_SOURCE="{BASH_SOURCE[0]}")

create_dir_structure_script='''
#! /bin/bash
mkdir -p {subpackage_dir}
mkdir -p {subpackage_dir}/modules
mkdir -p {subpackage_dir}/modules/datastructures
mkdir -p {subpackage_dir}/scripts
mkdir -p {subpackage_dir}/Train
mkdir -p {subpackage_dir}/example_data
'''.format(subpackage_dir=subpackage_dir)

datastructure_template='''



from DeepJetCore.TrainData import TrainData, fileTimeOut
import numpy 

class TrainData_example(TrainData):
    def __init__(self):
        TrainData.__init__(self)

        self.treename="tree" #input root tree name
        
        self.truthclasses=['isA','isB','isC'] #truth classes for classification
        
        self.weightbranchX='isA' #needs to be specified if weighter is used
        self.weightbranchY='isB' #needs to be specified if weighter is used
        
        #there is no need to resample/reweight
        self.weight=False
        self.remove=False
        #does not do anything in this configuration
        self.referenceclass='flatten'
        self.weight_binX = numpy.array([0,40000],dtype=float) 
        self.weight_binY = numpy.array([0,40000],dtype=float) 
        
        
        self.registerBranches(['']) #list of branches to be used 
        
        self.registerBranches(self.truthclasses)
        
        
        #call this at the end
        self.reduceTruth(None)
        
    
    def readFromRootFile(self,filename,TupleMeanStd, weighter):
    
        # this function defines how to convert the root ntuple to the training format
        # options are not yet described here
        
        import ROOT
        fileTimeOut(filename,120) #give eos a minute to recover
        rfile = ROOT.TFile(filename)
        tree = rfile.Get("tree")
        self.nsamples=tree.GetEntries()
        
        
        # user code, example works with the example 2D images in root format generated by make_example_data
        from DeepJetCore.preprocessing import read2DArray
        print(filename)
        feature_array = read2DArray(filename,"tree","image2d",self.nsamples,32,32)
        
        
        truth = self.read_truthclasses(filename)
        
        #notremoves=weighter.createNotRemoveIndices(Tuple)
        
        # this removes parts of the dataset for weighting the events
        #feature_array = feature_array[notremoves > 0]
                
        # call this in the end
        
        self.nsamples=len(feature_array)
        
        self.x=[feature_array] # list of feature numpy arrays
        self.y=[truth] # list of target numpy arrays (truth)
        self.w=[] # list of weight arrays. One for each truth target, not used



class TrainData_example_reg(TrainData):
    def __init__(self):
        TrainData.__init__(self)

        self.treename="tree" #input root tree name
        
        self.truthclasses=['isA','isB','isC'] #truth classes for classification
        
        self.weightbranchX='isA' #needs to be specified if weighter is used
        self.weightbranchY='isB' #needs to be specified if weighter is used
        
        #there is no need to resample/reweight
        self.weight=False
        self.remove=False
        #does not do anything in this configuration
        self.referenceclass='flatten'
        self.weight_binX = numpy.array([0,40000],dtype=float) 
        self.weight_binY = numpy.array([0,40000],dtype=float) 
        
        
        self.registerBranches(['']) #list of branches to be used 
        
        self.registerBranches(self.truthclasses)
        
        
        #call this at the end
        self.reduceTruth(None)
        
    
    def readFromRootFile(self,filename,TupleMeanStd, weighter):
    
        # this function defines how to convert the root ntuple to the training format
        # options are not yet described here
        
        import ROOT
        fileTimeOut(filename,120) #give eos a minute to recover
        rfile = ROOT.TFile(filename)
        tree = rfile.Get("tree")
        self.nsamples=tree.GetEntries()
        
        
        # user code, example works with the example 2D images in root format generated by make_example_data
        from DeepJetCore.preprocessing import read2DArray
        print(filename)
        feature_array = read2DArray(filename,"tree","image2d",self.nsamples,32,32)
        
        npy_array = self.readTreeFromRootToTuple(filename)
        reg_truth = npy_array['sigsum']
        
        #notremoves=weighter.createNotRemoveIndices(Tuple)
        
        # this removes parts of the dataset for weighting the events
        #feature_array = feature_array[notremoves > 0]
                
        # call this in the end
        
        self.nsamples=len(feature_array)
        
        self.x=[feature_array] # list of feature numpy arrays
        self.y=[reg_truth] # list of target numpy arrays (truth)
        self.w=[] # list of weight arrays. One for each truth target, not used


'''


training_template='''


from DeepJetCore.training.training_base import training_base
import keras
from keras.models import Model
from keras.layers import Dense, Conv2D, Flatten #etc

def my_model(Inputs,nclasses,nregressions,otheroption):
    
    x = Inputs[0] #this is the self.x list from the TrainData data structure
    x = Conv2D(8,(4,4),activation='relu', padding='same')(x)
    x = Conv2D(8,(4,4),activation='relu', padding='same')(x)
    x = Conv2D(8,(4,4),activation='relu', padding='same')(x)
    x = Conv2D(8,(4,4),strides=(2,2),activation='relu', padding='valid')(x)
    x = Conv2D(4,(4,4),strides=(2,2),activation='relu', padding='valid')(x)
    x = Flatten()(x)
    x = Dense(32, activation='relu')(x)
    
    x = Dense(nclasses)(x)
    
    predictions = [x]
    return Model(inputs=Inputs, outputs=predictions)

def my_regression_model(Inputs,nclasses,nregressions,otheroption):
    
    x = Inputs[0] #this is the self.x list from the TrainData data structure
    x = Conv2D(8,(4,4),activation='relu', padding='same')(x)
    x = Conv2D(8,(4,4),activation='relu', padding='same')(x)
    x = Conv2D(8,(4,4),activation='relu', padding='same')(x)
    x = Conv2D(8,(4,4),strides=(2,2),activation='relu', padding='valid')(x)
    x = Conv2D(4,(4,4),strides=(2,2),activation='relu', padding='valid')(x)
    x = Flatten()(x)
    x = Dense(32, activation='relu')(x)
    
    x = Dense(nregressions, activaton='None')(x)
    
    predictions = [x]
    return Model(inputs=Inputs, outputs=predictions)



train=training_base(testrun=False,resumeSilently=False,renewtokens=True)


if not train.modelSet(): # allows to resume a stopped/killed training. Only sets the model if it cannot be loaded from previous snapshot

    #for regression use the regression model
    train.setModel(my_model,otheroption=1)
    
    #for regression use a different loss, e.g. mean_squared_error
    train.compileModel(learningrate=0.003,
                   loss='categorical_crossentropy') 
                   
print(train.keras_model.summary())


model,history = train.trainModel(nepochs=10, 
                                 batchsize=500,
                                 checkperiod=1, # saves a checkpoint model every N epochs
                                 verbose=1)





'''
        
datastructures_init = '''
#Make it look like a package
from glob import glob
from os import environ
from os.path import basename, dirname
from pdb import set_trace

#gather all the files here
modules = [basename(i.replace('.py','')) for i in glob('%s/[A-Za-z]*.py' % dirname(__file__))]
__all__ = []
structure_list=[]
for module_name in modules:
    module = __import__(module_name, globals(), locals(), [module_name])
    for model_name in [i for i in dir(module) if 'TrainData' in i]:
        
        
        model = getattr(module, model_name)
        globals()[model_name] = model
        locals( )[model_name] = model
        __all__.append(model_name)
        structure_list.append(model_name)

'''
        
layers_template='''
# Define custom layers here and add them to the global_layers_list dict (important!)
global_layers_list = {}
'''
losses_template='''
# Define custom losses here and add them to the global_loss_list dict (important!)
global_loss_list = {}
'''

metrics_template='''
# Define custom metrics here and add them to the global_metrics_list dict (important!)
global_metrics_list = {}
'''


######## create the structure ########


os.system(create_dir_structure_script)
with  open(subpackage_dir+'/env.sh','w') as envfile:
    envfile.write(environment_file)
    
with  open(subpackage_dir+'/modules/datastructures/TrainData_example.py','w') as lfile:
    lfile.write(datastructure_template)
    
with  open(subpackage_dir+'/modules/datastructures/__init__.py','w') as lfile:
    lfile.write(datastructures_init)
    
with  open(subpackage_dir+'/Train/training_example.py','w') as lfile:
    lfile.write(training_template)
    
with  open(subpackage_dir+'/modules/Layers.py','w') as lfile:
    lfile.write(layers_template)
with  open(subpackage_dir+'/modules/Losses.py','w') as lfile:
    lfile.write(losses_template)
with  open(subpackage_dir+'/modules/Metrics.py','w') as lfile:
    lfile.write(metrics_template)
    
if args.nodata:
    exit()
print('creating example data... (10 training files, 1 test file, 1000 events each)')
os.system('cd '+subpackage_dir+'/example_data;  make_example_data  1000 10 1')


print('example data can be found in '+subpackage_dir+'/example_data.')
print('Before using the subpackage, please log out, log in again and then source the "env.sh" file in the subpackage directory (not in DeepJetCore).')
print('to convert to example TrainData format use:')
print('convertFromRoot.py -i '+subpackage_dir+'/example_data/train_files.txt -o <train output dir> -c TrainData_example')
print('to convert the test data use:')
print('convertFromRoot.py -i '+subpackage_dir+'/example_data/test_files.txt -o <test output dir> --testdatafor <train output dir>/dataCollection.dc ')
print('\nAn example to run the training can be found in '+subpackage_dir+'/Train/training_example.py')
print('It can be run with: \npython '+subpackage_dir+'/Train/training_example.py <train output dir>/dataCollection.dc <train output dir>')










