'''Rolling_split a order file and settle'''
'''Split orders in to obv windows and evaluation windows in order to try to predict evaluation window with obv window'''

'''pre_process.py'''
import sys
import argparse
import ConfigParser
import datetime   

from datetime import timedelta
from pyspark import SparkContext
from pyspark.sql import SparkSession

from pyspark.sql.types import *
from pyspark.sql.functions import udf
from pyspark.storagelevel import StorageLevel

from keystone.transformers.order_matcher import OrderMatcher
from keystone.data import schemas

from keystone.utils.spark_reader import load

def build_parser():
    parser = argparse.ArgumentParser(prog = 'rolling_split_and_settle.py', description = 'Rollingsplit orders and settle.')
    parser.add_argument('--master', dest = 'master', default = 'local[2]')
    parser.add_argument('--input', dest = 'input', required = True)
    parser.add_argument('--obswindow',dest = 'obswindow', required = True, type = int)
    parser.add_argument('--evawindow',dest = 'evawindow', required = True, type = int)
    parser.add_argument('--rollingstep',dest = 'rollingstep', default = 1, type = int)
    # parser.add_argument('--format', dest = 'format', default = 'tsv', choices = ['tsv', 'csv', 'pqt'])
    # parser.add_argument('--config', dest = 'config', required = True)
    parser.add_argument('--output', dest = 'output', required = True)
    parser.add_argument('--mode', dest = 'mode', default = 'ignore')
    return parser


def settle_window(orders):
    '''This method is designed for settle the orders within a given window (3 months or 1 month)
    and return the settled orders'''

    # orders is a spark dataframe generated by pre_processing.py
		


    # transform orders to user dataframe
    user_orders = orders.rdd.map(lambda order: (order['Trader'], schemas.order_to_tuple(order)))

    order_num_partitions = None
    
    user_orders = user_orders.aggregateByKey([],
                                        lambda orders, order: orders + [order],
                                        lambda orders1, orders2: orders1 + orders2,
                                        numPartitions=order_num_partitions)

    
    user_dataframe = spark.createDataFrame(user_orders, schema=schemas.orders_dataframe_schema)

    user_dataframe = user_dataframe.withColumn('LastOrders', udf(lambda order: [], ArrayType(schemas.order_schema))(user_dataframe['Orders']))
    user_dataframe = user_dataframe.withColumn('LastClosedPositions', udf(lambda order: [], ArrayType(schemas.closed_position_schema))(user_dataframe['Orders']))
    user_dataframe = user_dataframe.withColumn('LastOpenPositions', udf(lambda order: [], ArrayType(schemas.open_position_schema))(user_dataframe['Orders']))
    user_dataframe = user_dataframe.withColumn('LastAccountDaily', udf(lambda order: None, schemas.account_daily_schema)(user_dataframe['Orders']))
    user_dataframe = user_dataframe.withColumn('LastAccountDailies', udf(lambda order: [], ArrayType(schemas.account_daily_schema))(user_dataframe['Orders']))
    dataframe = user_dataframe

    settlements = dataframe.persist(storageLevel=StorageLevel.MEMORY_AND_DISK)

    
    # run settlement
    order_matcher = OrderMatcher()
    settlements = order_matcher.transform(settlements)

    settlements = settlements.drop('LastAccountDaily')
    settlements = settlements.drop('LastAccountDailies')

    # write user dataframe
    # settlements.printSchema()
    
    return settlements

if __name__ == '__main__':
    # load command line args and config
    cmd_args_parser = build_parser()
    args = cmd_args_parser.parse_args()
    #config = load_config(args.config)

    # initialize reader
    sc = SparkContext(appName='Rolling Split')
    spark = SparkSession(sc)
    reader = spark.read
    
    # transform order schema
    orders = load(reader, args.input, 'rolling_split')
    orders = orders.sort(orders.Timestamp)
    orders.show()
    start_of_trade = orders.select('Timestamp').rdd.min()[0]
    start_of_trade = start_of_trade.replace(hour=0, minute=0, second=0, microsecond=0)
    end_of_trade = orders.select('Timestamp').rdd.max()[0]
    end_of_trade = end_of_trade.replace(hour=0, minute=0, second=0, microsecond=0)
    print('executed')

    rolling_window_obv = timedelta(days = args.obswindow)
    rolling_window_evl = timedelta(days = args.evawindow)
    rolling_step = timedelta(days = args.rollingstep)
    observe_start = start_of_trade
    observe_end = observe_start + rolling_window_obv
    eval_start = observe_end
    eval_end = observe_end + rolling_window_evl
    #print '\n\n\n\n\n'
    #print rolling_window_obv, rolling_window_evl, observe_start, observe_end
    #print eval_start, eval_end
    #print '\n\n\n\n\n'
    observation = orders.filter((orders.Timestamp >= observe_start) & (orders.Timestamp < observe_end))
    evaluation = orders.filter((orders.Timestamp >= eval_start) & (orders.Timestamp <eval_end))
    if observation.rdd.isEmpty() or evaluation.rdd.isEmpty():
        print observation.rdd.isEmpty(), evaluation.rdd.isEmpty()
        raise ValueError('obsevation or evaluation is empty. Check parameters')

    settle_obs = settle_window(observation)
    settle_eva = settle_window(evaluation)


    settle_eva = settle_window(evaluation) #rename columns to avoid name conflicts
    settle_eva = settle_eva.withColumnRenamed('TodayOrders', 'TodayOrders_eva')
    settle_eva = settle_eva.withColumnRenamed('TodayClosedPositions', 'TodayClosedPositions_eva')
    settle_eva = settle_eva.withColumnRenamed('OpenPositions', 'OpenPositions_eva')
    settle_eva = settle_eva.withColumnRenamed('ClosedPositions', 'ClosedPositions_eva')
    settle_eva = settle_eva.withColumnRenamed('Orders', 'Orders_eva')
	
	# Join the two DataFrame so that the joined_settle has a single trader's settlements on each row, left part is the obv period(3 month)
	#  and the right part is the evaluation period(1 month)
    rolling_settlements = settle_obs.join(settle_eva, on=['Trader'], how='inner')

    while eval_end < end_of_trade:

	
    	observe_start += rolling_step                    # move the starting point of a observation period by rolling_step
    	observe_end = observe_start + rolling_window_obv # observe_end - observe_start = length of observation period
		
    	eval_start = observe_end
    	eval_end = observe_end + rolling_window_evl

    	if eval_end > end_of_trade:
            print '\n\n\n\n\n'
            #raise ValueError('END OF While!!!')
            break
  
    	observation = orders.filter((orders.Timestamp >= observe_start) & (orders.Timestamp < observe_end))
    	evaluation = orders.filter((orders.Timestamp >= eval_start) & (orders.Timestamp <eval_end))
	print '\n\n\n'
        print '\n a new Observation Window\n'
	observation.show()


    	if observation.rdd.isEmpty() or evaluation.rdd.isEmpty(): # if no trade, skip this window
            continue

		# settle_obs is the settlement for the observe period, e.g. 3 month
		# settle_eva is the settlement for the evaluation period following the observation period, for example, the following 1 month
    	settle_obs = settle_window(observation)
    	settle_eva = settle_window(evaluation)

		# rename the columns of settle_eva to avoid column name conflicts with settle_obs
    	settle_eva = settle_eva.withColumnRenamed('TodayOrders', 'TodayOrders_eva')
    	settle_eva = settle_eva.withColumnRenamed('TodayClosedPositions', 'TodayClosedPositions_eva')
    	settle_eva = settle_eva.withColumnRenamed('OpenPositions', 'OpenPositions_eva')
    	settle_eva = settle_eva.withColumnRenamed('ClosedPositions', 'ClosedPositions_eva')
    	settle_eva = settle_eva.withColumnRenamed('Orders', 'Orders_eva')

        settle_eva.printSchema()


    	joined_settlements = settle_obs.join(settle_eva, on=['Trader'], how='inner')

	# append the current observation & evaluation pair to the rolling_settlements 
    	rolling_settlements = rolling_settlements.unionAll(joined_settlements)

    #rolling_settlements.show()
    rolling_settlements.printSchema()
    #rolling_settlements.write.save(args.output, mode=args.mode)
    # write orders as parquet
    # orders.write.save(args.output, mode=args.mode)



