import time
import pandas as pd
import numpy as np
import datetime as dt
from collections import OrderedDict
from copy import copy
import warnings
import matplotlib.pyplot as plt
import seaborn as sns
from pprint import pprint
%matplotlib inline

"""
Explanation: Import necessary modules
End of explanation
"""


project_dir = r'/Users/hudson/Code/marketModel/'

"""
Explanation: Filepath management
End of explanation
"""


stock_data = pd.read_hdf(project_dir + 'data/stock_data/raw_stock_data.hdf', 'table')
symbols = stock_data.reset_index().ticker.unique()
pprint(symbols)

symbol = np.random.choice(symbols)
print 'symbol: ' + symbol
stock_data_vis = stock_data.loc[symbol]
print stock_data_vis.head()
print stock_data_vis.describe()
print stock_data.reset_index().loc[:,('ticker', 'timestamp')].groupby('ticker').agg([len, np.min, np.max])

"""
Explanation: Load the data from the hdf store
End of explanation
"""


# First keep the time index so we can see the time frame
stock_data_vis.close.plot(label=symbol)

plt.legend(bbox_to_anchor=(1.25, .5))
plt.tight_layout()
plt.ylabel("Close ($)")
sns.despine()

# Now drop the time index so we can see the actual stock movement
stock_data_vis.reset_index().close.plot(label=symbol)
stock_data_vis.reset_index().close.rolling(20).mean().plot(label='20 Min. Moving Avg.')

plt.legend(bbox_to_anchor=(1.25, .5))
plt.tight_layout()
plt.ylabel("Close ($)")
sns.despine()

"""
Explanation: Visualize the data
End of explanation
"""


# First reshape the index and group by ticker
stock_data_final = stock_data.reset_index(level=0)
grouped_stock_data = stock_data_final.groupby('ticker')

## Before evaluating goals, we need to fill in the missing time steps.
## These missing time steps have zero trading volume.
## So, all of the prices for these steps are equal to the last
## closing price. All of the volumes are equal to zero.

stock_data_with_all_minutes = []
for name, group in grouped_stock_data:
    
    # Create a dataframe of all the times
    min_time, max_time = group.index.min(), group.index.max()
    timeDiff = (max_time - min_time).components
    numMinutes = timeDiff.days*24*60 + timeDiff.hours*60 + timeDiff.minutes

    #alltimesIdx = pd.DatetimeIndex(start=min_time, freq=pd.tseries.offsets.Minute(1), periods=numMinutes)
    alltimesIdx = pd.DatetimeIndex(start=min_time, freq=pd.Timedelta(minutes=1), periods=numMinutes)
    alltimes = pd.DataFrame(index=alltimesIdx)

    # Drop minutes outside of 9:30am - 4:00pm est
    alltimes = alltimes.between_time('09:30','16:00')
  
    # Join on the original dataframe
    alltimes_group = alltimes.join(group)
    
    # Forward fill the NaN closing prices 
    alltimes_group.loc[:,('ticker', 'close')] = alltimes_group.loc[:,('ticker', 'close')].\
                                                    fillna(method='ffill', axis=0)
    
    # Assign all price variables to the close price
    alltimes_group.loc[:,'open':'close'] = alltimes_group.loc[:,'open':'close'].\
                                                fillna(method='bfill', axis=1)
    
    # Assign all NaN volumes to zero
    alltimes_group.loc[:, 'volume'] = alltimes_group.loc[:, 'volume'].fillna(value=0)
    
    stock_data_with_all_minutes.append(alltimes_group)
    
stock_data_with_all_minutes = pd.concat(stock_data_with_all_minutes)
stock_data_with_all_minutes.index.name = 'timestamp'

stock_data_with_all_minutes.reset_index().loc[:,('ticker', 'timestamp')].groupby('ticker').agg([len, min, max])

"""
Explanation: Adding in missing times (zero volume minutes)
Before evaluating goals, we need to fill in the missing time steps.
These missing time steps have zero trading volume.
So, all of the prices for these steps are equal to the last
closing price. All of the volumes are equal to zero.
End of explanation
"""


def get_min_max(data, starttime, endtime):
    """
    This function takes data for a specific ticker and returns the min and max prices.
    """
    subdata = data.loc[starttime:endtime]
    
    return (subdata.low.min(), subdata.high.max())
    
def is_goal_met(data, timestep, goal_time_from_step, goal_duration, goal_raise_frac = 0.1, goal_drop_frac=0.1):
    """
    This function takes data for a specific ticker, a time index for that ticker, goal parameters, and 
    returns a boolean indicating whether or not the goal is satisfied for that timestep.
    """
    
    # Assign a status message to record various special cases
    statusMessage = ""
    
    #Convert time variables to appropriate numpyt date types
    td_goal_time_from_step = np.timedelta64(goal_time_from_step, 'm')
    td_goal_duration = np.timedelta64(goal_duration, 'm')
    
    # Calculate the start and end times of the goal time window
    goal_starttime = np.datetime64(timestep + td_goal_time_from_step)
    goal_endtime = np.datetime64(goal_starttime + td_goal_duration)
    
    if goal_endtime > np.datetime64(data.index.max()):
        statusMessage = "Goal time window end time lies beyond available data."
    
    # Get the data for goal checking in that time window
    subdata = data.loc[goal_starttime:goal_endtime]
    
    # Get the minimum and maximum prices for the goal time window
    min_price, max_price = get_min_max(data, goal_starttime, goal_endtime)
    if np.isnan(min_price) | np.isnan(max_price):
        # Zero trading volume in time window. Get last prices.
        most_recent_time_with_transactions = np.max(data.loc[:goal_starttime].index)
        if most_recent_time_with_transactions == timestep:
            statusMessage = statusMessage + " Zero trading volume between current timestep and goal time window end."
            return {'timestamp': timestep,
                    'goal_met': False,
                    'raise_goal_met': False,
                    'drop_goal_met': True,
                    'statusMessage': statusMessage}
        else:
            min_price, max_price = data.loc[timestep, 'low'], data.loc[timestep, 'high']
        
        
    # Determine if goals were met
    # TODO: is this the right reference for the 'current price'?
    current_price = np.mean(data.loc[timestep, ['high', 'close']])
    
    # Is raise goal met? Return true if max price at least (1+goal_raise_frac) * current_price
    is_raise_goal_met = max_price >= (1+goal_raise_frac) * current_price
    
    # Is drop goal met? Return true if min price at least (1-goal_drop_frac) * current_price
    is_drop_goal_met = min_price >= (1-goal_drop_frac) * current_price
    
    # Return dict containing raise and drop goals and product for convenience
    return {'timestamp': timestep,
        'goal_met': is_raise_goal_met * is_drop_goal_met, 
            'raise_goal_met': is_raise_goal_met, 
            'drop_goal_met': is_drop_goal_met,
             'statusMessage': statusMessage}

# test get_min_max
get_min_max(stock_data_vis, '2017-08-28', '2017-08-30')

# test is_goal_met
random_time_index = np.random.choice(stock_data_vis.index.values)
print "Random time: " + str(random_time_index)
%timeit is_goal_met(stock_data_vis, random_time_index, 0, 1000)

"""
Explanation: Logic for getting goal information
End of explanation
"""


# first define a function that tags for one ticker
def get_tagged_stock_data(data, 
                          ticker, 
                          goal_time_from_step, 
                          goal_duration, 
                          goal_raise_frac = 0.1, 
                          goal_drop_frac=0.1):
    # Loop over the timestamps building a dictionary of the tagging information
    tagged_stock_data = []
    for timestep in data.index:
        goal_dict = is_goal_met(data, timestep, goal_time_from_step, goal_duration, goal_raise_frac, goal_drop_frac)
        tagged_stock_data.append(goal_dict)
    # Convert to pandas and return
    return pd.DataFrame(tagged_stock_data).set_index('timestamp')

start_delay = 0 # minutes
duration = 120 # minutes (next half hour)
raise_fraction = 0.05
drop_fraction = 0.05
list_tagged_tickers = []
for i, (symbol, group) in enumerate(stock_data_with_all_minutes.groupby('ticker')):
    print "Progress: {} of {} tickers. Current ticker: {}".format(i, len(symbols), symbol)
    
    # get the tag data for this symbol
    tag_data = get_tagged_stock_data(group, symbol, start_delay, duration, raise_fraction, drop_fraction)
    
    # join tag data back onto the group data
    merged_data = group.join(tag_data)
    
    # Append to the list of tagged data
    list_tagged_tickers.append(merged_data)
    
# Concatenate all the groups
all_tagged_data = pd.concat(list_tagged_tickers)

#print all_tagged_data.goal_met.value_counts()
#print all_tagged_data.statusMessage.value_counts()
print all_tagged_data.groupby('ticker').agg({'goal_met': lambda x: x[x].shape[0],
'raise_goal_met': lambda x: x[x].shape[0],
'drop_goal_met': lambda x: x[x].shape[0]}).agg(np.sum)

"""
Explanation: Get goal tags for all data
End of explanation
"""


all_tagged_data.to_hdf(project_dir + 'data/stock_data/tagged_stock_data.hdf', 'table')

"""
Explanation: Save data to hdf
End of explanation
"""


