from zipline.api import (set_slippage, 
                         set_commission, 
                         set_benchmark, 
                         attach_pipeline, 
                         symbol, 
                         pipeline_output,
                         record, 
                         order,
                         order_target
                        )
from zipline.pipeline.filters import StaticSids
from zipline.finance import slippage, commission
from zipline import run_algorithm
from zipline.pipeline import CustomFactor, Pipeline
from zipline.pipeline.data import EquityPricing
from zipline.pipeline.factors import ExponentialWeightedMovingAverage

def make_pipeline():
    ema = ExponentialWeightedMovingAverage(inputs = [EquityPricing.close],window_length = 7,decay_rate = 1/7)
    high = NdaysMaxHigh(inputs = [EquityPricing.close], window_length = 8) # window_length 設定為 8，因為 factor 會包含當日價格。
    low = NdaysMinLow(inputs = [EquityPricing.close], window_length = 8)
    close = EquityPricing.close.latest
    return Pipeline(
        columns = {
            'ema':ema,
            'highesthigh':high,
            'lowestlow':low,
            'latest':close
        }
    )
class NdaysMaxHigh(CustomFactor):
    def compute(self, today, assets, out, data):
        out[:] = np.nanmax(data[:-2], axis=0)
class NdaysMinLow(CustomFactor):
    def compute(self, today, assets, out, data):
        out[:] = np.nanmin(data[:-2], axis=0)

def initialize(context):
    set_slippage(slippage.VolumeShareSlippage())
    set_commission(commission.PerShare(cost=0.00285))
    set_benchmark(symbol('2330'))
    attach_pipeline(make_pipeline(), 'mystrategy')

def handle_data(context, data):
    
    pipe = pipeline_output('mystrategy')
    
    for i in pipe.index:
        ema = pipe.loc[i, 'ema']
        highesthigh = pipe.loc[i, 'highesthigh']
        lowestlow = pipe.loc[i, 'lowestlow']
        close = pipe.loc[i, 'latest']
        bias = close - ema
        residual_position = context.portfolio.positions[i].amount # 當日該資產的股數
        condition1 = (close > highesthigh) and (bias > 0) and (residual_position > 0) # 賣出訊號
        condition2 = (close < lowestlow) and (bias < 0) # 買入訊號
        
        record( # 用以紀錄以下資訊至最終產出的 result 表格中
            con1 = condition1,
            con2 = condition2,
            price = close,
            ema = ema,
            bias = bias,
            highesthigh = highesthigh,
            lowestlow = lowestlow
        )
    
        if condition1:
            order_target(i, 0)
        elif condition2:
            order(i, 10)
        else:
            pass

import matplotlib.pyplot as plt
def analyze(context, perf):
    fig = plt.figure()
    ax1 = fig.add_subplot(211)
    perf.portfolio_value.plot(ax=ax1)
    ax1.set_ylabel("Portfolio value (NTD)")
    ax2 = fig.add_subplot(212)
    ax2.set_ylabel("Price (NTD)")
    perf.price.plot(ax=ax2)
    ax2.plot( # 繪製買入訊號
        perf.index[perf.con2],
        perf.loc[perf.con2, 'price'],
        '^',
        markersize=5,
        color='red'
    )
    ax2.plot( # 繪製賣出訊號
        perf.index[perf.con1],
        perf.loc[perf.con1, 'price'],
        'v',
        markersize=5,
        color='green'
    )
    plt.legend(loc=0)
    plt.gcf().set_size_inches(18,8)
    plt.show()

import pyfolio as pf 
returns, positions, transactions = pf.utils.extract_rets_pos_txn_from_zipline(results)

benchmark_rets = results['benchmark_return'] 
from pyfolio.utils import extract_rets_pos_txn_from_zipline
returns, positions, transactions = extract_rets_pos_txn_from_zipline(results)