# Pipeline imports
from quantopian.pipeline import CustomFactor, Pipeline
from quantopian.pipeline.filters import QTradableStocksUS, StaticAssets
from quantopian.pipeline.factors import DailyReturns

# Import run_pipeline method
from quantopian.research import run_pipeline

# Import optimize
import quantopian.experimental.optimize as opt

# Pipeline definition 
# Make a pipeline with defined percent of longs and shorts
def  make_pipeline(universe_size=100, short_percent=50): 
    factor_universe = QTradableStocksUS()
    
    # Some random data. Universe with highest returns.
    my_universe = DailyReturns(mask=factor_universe).top(universe_size)
    
    # Filter for specific assets to pass to CompareAssets factor
    # Could be the result of some calculation but just using a static filter for testing
    shorts = DailyReturns(mask=my_universe).percentile_between(0, short_percent)
    longs = my_universe & ~shorts
        
    return Pipeline(
        columns={
            'shorts': shorts,
            'longs': longs,
                },
        screen=my_universe
    )

def get_weights(universe_size=100, short_percent=50, date='2019-01-02'):
    df = run_pipeline(make_pipeline(universe_size=universe_size, 
                                    short_percent=short_percent), date, date)
    
    # Create a column for weight - equal weight longs and shorts
    weight = 1.0 / df.index.size
    df.loc[df.longs, 'weights'] = weight
    df.loc[df.shorts, 'weights'] = -weight    
    
    return df.xs(date, level=0).weights

# Lets run optimize to limit our short exposure
# Create a set of weights with 90% short
orig_weights = get_weights(universe_size=100, short_percent=90)

# One issue is that the optimize constraint defines 'net exposure' and not 'short weight'
# So, we need to do a bit of conversion. Start with our max desired short weight
max_percent_shorts = .8

# This implies the following max short exposure 
max_short_exposure = (max_percent_shorts * -1.0) + (1.0-max_percent_shorts)

weights = orig_weights
loop = 0
while True:
    # Now calculate an optimized portfolio using a net_exposure_constraint
    objective = opt.TargetWeights(weights)
    net_exposure_constraint = opt.NetExposure(max_short_exposure, 1.0)
    opt_weights = opt.calculate_optimal_portfolio(objective, constraints=[net_exposure_constraint])
    gross_exposure = opt_weights.abs().sum()
    weights = opt_weights / gross_exposure
    loop += 1
    if gross_exposure > .99:
        break


# Let's get some information about our results
short_percent = opt_weights.where(opt_weights<0).sum()
long_percent = opt_weights.where(opt_weights>0).sum()
net_exposure = short_percent + long_percent
gross_exposure = abs(short_percent) + abs(long_percent)

display('number of iterations: {}'.format(loop))
display('gross exposure: {}'.format(gross_exposure))
display('net exposure: {}'.format(net_exposure))
display('short percent: {}'.format(short_percent))
display(opt_weights)

# IT WORKS!
# But notice one had to iterate 6 times to get leverage above .99