text: import wrds
  import datetime
  import pandas as pd
  import trading_calendars as tc
  import pytz
  import numpy as np
  from stockstats import StockDataFrame as Sdf
  pd.options.mode.chained_assignment = None 
  
  class WrdsEngineer():
      def __init__(self,if_offline=False):
          if not if_offline:
              self.db = wrds.Connection()
  
      def data_fetch_ohlcv(self, start, end, stock_list, time_interval, if_save_tempfile=False,
                           filter_shares=0):
          
      
      
          def get_trading_days(start, end):
              nyse = tc.get_calendar('NYSE')
              df = nyse.sessions_in_range(pd.Timestamp(start,tz=pytz.UTC),
                                          pd.Timestamp(end,tz=pytz.UTC))
              trading_days = []
              for day in df:
                  trading_days.append(str(day)[:10])
          
              return trading_days
      
          
          def data_fetch_wrds(date='2021-05-01',
                              stock_set=('AAPL'), time_interval=60):
              #start_date, end_date should be in the same year
              current_date = datetime.datetime.strptime(date, '%Y-%m-%d')
              lib = 'taqm_' + str(current_date.year) #taqm_2021
              table = 'ctm_'+ current_date.strftime('%Y%m%d') #ctm_20210501
      
              parm = {'syms' : stock_set, 'num_shares': filter_shares}
              try:
                  data = self.db.raw_sql("select * from " + lib + '.'+ table + 
                                    " where sym_root in %(syms)s and time_m between '9:30:00' and '16:00:00' and size > %(num_shares)s",
                                    params = parm)
                  if_empty = False
                  return data, if_empty
              except:
                  print('Data for date: ' + date + ' error')
                  if_empty = True
                  return None, if_empty
              
          dates = get_trading_days(start, end)
          print('Trading days: ')
          print(dates)
          first_time = True
          empty = True
          stock_set = tuple(stock_list)
          for i in dates:            
              x = data_fetch_wrds(i, stock_set, 
                                    time_interval)
              
              if not x[1]:
                  empty = False
                  dataset = x[0]
                  dataset = self.preprocess_to_ohlcv(dataset, time_interval = (str(time_interval)+'S'))
                  if first_time:
                      print('Data for date: ' + i + ' finished')
                      temp = dataset
                      first_time = False
                      if if_save_tempfile:
                          temp.to_csv('./temp.csv')
                  else:
                      print('Data for date: ' + i + ' finished')
                      temp = pd.concat([temp,dataset])
                      if if_save_tempfile:
                          temp.to_csv('./temp.csv')
          if empty:
              raise ValueError('Empty Data under input parameters!')
          else:
              result = temp
              result = result.sort_values(by=['time','tic'])
              result = result.reset_index(drop=True)
              return result
      
      def preprocess_to_ohlcv(self, df, time_interval='60S'):
          df = df[['date','time_m','sym_root','size','price']]
          tic_list = np.unique(df['sym_root'].values)
          final_df = None
          first_time = True
          for i in range(len(tic_list)):
              tic = tic_list[i]
              time_list = []
              temp_df = df[df['sym_root']==tic]
              for i in range(0,temp_df.shape[0]):
                  date = temp_df['date'].iloc[i]
                  time_m = temp_df['time_m'].iloc[i]
                  time = str(date) + ' ' + str(time_m)
                  try:
                      time = datetime.datetime.strptime(time, '%Y-%m-%d %H:%M:%S.%f')
                  except:
                      time = datetime.datetime.strptime(time, '%Y-%m-%d %H:%M:%S')
                  time_list.append(time)
              temp_df['time'] = time_list
              temp_df = temp_df.set_index('time')
              data_ohlc = temp_df['price'].resample(time_interval).ohlc()
              data_v = temp_df['size'].resample(time_interval).agg({'size':'sum'})
              volume = data_v['size'].values
              data_ohlc['volume'] = volume
              data_ohlc['tic'] = tic
              if first_time:
                  final_df = data_ohlc.reset_index()
                  first_time = False
              else:
                  final_df = final_df.append(data_ohlc.reset_index(),ignore_index=True)
          return final_df
      
      def data_clean(self, df):
          df = df[['time', 'open', 'high', 'low', 'close', 'volume', 'tic']]
          # remove 16:00 data
          tic_list = np.unique(df['tic'].values)
          ary = df.values
          rows_1600 = []
          for i in range(ary.shape[0]):
              row = ary[i]
              time = row[0]
              if str(time)[-8:] == '16:00:00':
                  rows_1600.append(i)
  
          df = df.drop(rows_1600)
          df = df.sort_values(by=['tic','time'])
          
          #check missing rows
          tic_dic = {}
          for tic in tic_list:
              tic_dic[tic] = [0,0]
          ary = df.values
          for i in range(ary.shape[0]):
              row = ary[i]
              volume = row[5]
              tic = row[6]
              if volume != 0:
                  tic_dic[tic][0] += 1
              tic_dic[tic][1] += 1
          constant = np.unique(df['time'].values).shape[0]
          nan_tics = []
          for tic in tic_dic:
              if tic_dic[tic][1] != constant:
                  nan_tics.append(tic)
          #fill missing rows
          normal_time = np.unique(df['time'].values) 
          
          df2 = df.copy()
          for tic in nan_tics:
              tic_time = df[df['tic'] == tic]['time'].values
              missing_time = []
              for i in normal_time:
                  if i not in tic_time:
                      missing_time.append(i)
              for time in missing_time:
                  temp_df = pd.DataFrame([[time,np.nan,np.nan,np.nan,np.nan,0,tic]], 
                                         columns=['time', 'open', 'high', 'low', 'close', 'volume', 'tic'])
                  df2=df2.append(temp_df,ignore_index=True)
          
          #fill nan data
          df = df2.sort_values(by=['tic','time'])
          for i in range(df.shape[0]):
              if float(df.iloc[i]['volume']) == 0:
                  previous_close = df.iloc[i-1]['close']
                  if str(previous_close) == 'nan':
                      raise ValueError('Error nan price')
                  df.iloc[i,1] = previous_close
                  df.iloc[i,2]= previous_close
                  df.iloc[i,3]= previous_close
                  df.iloc[i,4]= previous_close
          #check if nan 
          ary = df[['open','high','low','close','volume']].values
          assert np.isnan(np.min(ary)) == False
          #final preprocess
          df = df[['time','open','high','low','close','volume','tic']]
          df = df.reset_index(drop=True)
          print('Data clean finished')
          return df
      
      def add_technical_indicators(self, df, tech_indicator_list = [
              'macd', 'boll_ub', 'boll_lb', 'rsi_30', 'dx_30',
              'close_30_sma', 'close_60_sma']):
          df = df.rename(columns={'time':'date'})
          df = df.copy()
          df = df.sort_values(by=['tic', 'date'])
          stock = Sdf.retype(df.copy())
          unique_ticker = stock.tic.unique()
          tech_indicator_list = tech_indicator_list
          
          for indicator in tech_indicator_list:
              indicator_df = pd.DataFrame()
              for i in range(len(unique_ticker)):
                  # print(unique_ticker[i], i)
                  temp_indicator = stock[stock.tic == unique_ticker[i]][indicator]
                  temp_indicator = pd.DataFrame(temp_indicator)
                  temp_indicator['tic'] = unique_ticker[i]
                  # print(len(df[df.tic == unique_ticker[i]]['date'].to_list()))
                  temp_indicator['date'] = df[df.tic == unique_ticker[i]]['date'].to_list()
                  indicator_df = indicator_df.append(
                      temp_indicator, ignore_index=True
                  )
              df = df.merge(indicator_df[['tic', 'date', indicator]], on=['tic', 'date'], how='left')
          df = df.sort_values(by=['date', 'tic'])
          print('Succesfully add technical indicators')
          return df
      
      def df_to_ary(self,df,tech_indicator_list):
          columns = ['open','high','low','close','volume']
          columns.extend(tech_indicator_list)
          N = df.shape[0]
          stocks_num = len(np.unique(df['tic'].values))
          minutes = int(N/stocks_num)
          if_first_time = True
          for i in range(0,minutes):
              temp_ary = df.iloc[stocks_num * i:stocks_num*(i+1)][columns].values
              temp_ary = temp_ary.flatten()
              if if_first_time:
                  ary = temp_ary
                  if_first_time = False
              else:
                  ary = np.vstack((ary,temp_ary))
          print('Successfully transformed into array')
          return ary
          
          
      
              
  