import pandas as pd
import numpy as np
from pandas_datareader import data

start_date = '2014-01-01'
end_date = '2018-01-01'
SRC_DATA_FILENAME = 'goog_data.pkl'

try:
    goog_data = pd.read_pickle(SRC_DATA_FILENAME)
    print('File data found...reading GOOG data')
except FileNotFoundError:
    print('File not found...downloading the GOOG data')
    goog_data = data.DataReader('GOOG', 'yahoo', start_date, end_date)
    goog_data.to_pickle(SRC_DATA_FILENAME)

goog_data_signal = pd.DataFrame(index=goog_data.index)
goog_data_signal['price'] = goog_data['Adj Close']


import numpy as np
def trading_support_resistance(data, bin_width=20):

    data['sup_tolerance'] = pd.Series(np.zeros(len(data))) # tolerance of support line
    data['res_tolerance'] = pd.Series(np.zeros(len(data))) # tolerance of resistance line

    data['sup_count'] = pd.Series(np.zeros(len(data))) # number of hitting support line
    data['res_count'] = pd.Series(np.zeros(len(data)))

    data['sup'] = pd.Series(np.zeros(len(data))) # support line value (constant)
    data['res'] = pd.Series(np.zeros(len(data)))

    data['position'] = pd.Series(np.zeros(len(data)))
    data['signal'] = pd.Series(np.zeros(len(data)))

    in_support = 0
    in_resistance = 0

    for x in range((bin_width-1) + bin_width, len(data)):
        data_section = data[x-bin_width : x+1] # get data within rolling window
        
        support_level = min(data_section['price'])
        resistance_level = max(data_section['price'])
        range_level = resistance_level - support_level

        data['res'][x] = resistance_level
        data['sup'][x] = support_level

        data['sup_tolerance'][x] = support_level + 0.2*range_level # 20% tolorance
        data['res_tolerance'][x] = resistance_level - 0.2*range_level

        if (data['price'][x] >= data['res_tolerance'][x]) and (data['price'][x] <= data['res'][x]):
            # if price is within resistance tolerance region
            in_resistance += 1
            data['res_count'][x] = in_resistance
        elif (data['price'][x] <= data['sup_tolerance'][x]) and (data['price'][x] >= data['sup'][x]):
            # price is within support tolerance region
            in_support += 1
            data['sup_count'][x] = in_support
        else:
            # if not within any region, clear count
            in_support = 0
            in_resistance = 0

        if in_resistance > 2:
            # If enter resistance region twice
            data['signal'][x] = 1
        elif in_support > 2:
            data['signal'][x] = 0
        else:
            data['signal'][x] = data['signal'][x-1]
    
    data['position'] = data['signal'].diff()

trading_support_resistance(goog_data_signal)

print(goog_data_signal)

import matplotlib.pyplot as plt

fig = plt.figure()
ax1 = fig.add_subplot(111, ylabel='Google price in $')
goog_data_signal['sup'].plot(ax=ax1, color='g', lw=2.)
goog_data_signal['res'].plot(ax=ax1, color='b', lw=2.)
goog_data_signal['price'].plot(ax=ax1, color='r', lw=2.)

ax1.plot(goog_data_signal.loc[goog_data_signal.position == 1.0].index,
        goog_data_signal.price[goog_data_signal.position == 1.0], '^', markersize=7, color='k', label='buy')
ax1.plot(goog_data_signal.loc[goog_data_signal.position  == -1.0].index,
        goog_data_signal.price[goog_data_signal.position == -1.0], 'v', markersize=7, color='k', label='sell')

plt.legend()
plt.savefig("supportresistanceline_improved")
plt.show()