from AlgorithmImports import *

class ShortAvailabilityDataAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(2021, 1, 1)
        self.set_end_date(2021, 7, 1)
        self.set_cash(1000)
        # Seed the security price as the last known price, such that the price data is immediately available at initial rebalance
        self.set_security_initializer(MySecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        self.equity = self.add_equity("GME")

        # Set up daily rebalance scheduled event, since shortable quantity is updated daily
        self.schedule.on(
            self.date_rules.every_day(self.equity.symbol),
            self.time_rules.after_market_open(self.equity.symbol, 10),
            self.rebalance)

    def rebalance(self) -> None:
        symbol = self.equity.symbol
        
        # You can obtain the shortable quantity to decide the submission of a short order
        shortable_quantity = self.equity.shortable_provider.shortable_quantity(symbol, self.time)
        if not shortable_quantity:
            shortable_quantity = 0
        # Fee and rebate rate is also available, such that you can calculate the expected return and decide if the margin is worthwhile
        self.plot('Total Shortable Quantity', symbol, shortable_quantity)
        self.plot('Borrowing Cost', "Fee Rate", self.equity.shortable_provider.fee_rate(symbol, self.time))
        self.plot('Borrowing Cost', "Rebate Rate", self.equity.shortable_provider.rebate_rate(symbol, self.time))

        # Test whether we can short the desired quantity
        quantity = self.calculate_order_quantity(symbol, -1)
        if self.shortable(symbol, quantity):
            self.market_order(symbol, quantity)

    def on_margin_call_warning(self) -> None:
        self.liquidate()

class MySecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(self, brokerage_model: IBrokerageModel, security_seeder: ISecuritySeeder) -> None:
        super().__init__(brokerage_model, security_seeder)

    def initialize(self, security: Security) -> None:
        super().initialize(security)
        # Set the shortable provider as your broker for accurate short reality modeling
        security.set_shortable_provider(InteractiveBrokersShortableProvider())

using QuantConnect.Data.Shortable;
namespace QuantConnect
{
    public class ShortAvailabilityDataAlgorithm : QCAlgorithm
    {
        private Equity _equity;
        
        public override void Initialize()
        {
            SetStartDate(2021, 1, 1);
            SetEndDate(2021, 7, 1);
            SetCash(1000);
            // Seed the security price as the last known price, such that the price data is immediately available at initial rebalance
            SetSecurityInitializer(new MySecurityInitializer(BrokerageModel, new FuncSecuritySeeder(GetLastKnownPrices)));
            
            _equity = AddEquity("GME");
            
            // Set up daily rebalance scheduled event, since shortable quantity is updated daily
            Schedule.On(
                DateRules.EveryDay(_equity.Symbol),
                TimeRules.AfterMarketOpen(_equity.Symbol, 10),
                Rebalance);
        }
        
        public void Rebalance()
        {
            var symbol = _equity.Symbol;

            // You can obtain the shortable quantity to decide the submission of a short order
            // Fee and rebate rate is also available, such that you can calculate the expected return and decide if the margin is worthwhile
            Plot("Total Shortable Quantity", symbol, _equity.ShortableProvider.ShortableQuantity(symbol, Time) ?? 0m);
            Plot("Borrowing Cost", "Fee Rate", _equity.ShortableProvider.FeeRate(symbol, Time));
            Plot("Borrowing Cost", "Rebate Rate", _equity.ShortableProvider.RebateRate(symbol, Time));

            // Test whether we can short the desired quantity
            var quantity = CalculateOrderQuantity(symbol, -1m);
            if (Shortable(symbol, quantity))
            {
                MarketOrder(symbol, quantity);
            }
        }

        public override void OnMarginCallWarning()
        {
            Liquidate();
        }

        class MySecurityInitializer : BrokerageModelSecurityInitializer
        {
            public MySecurityInitializer(IBrokerageModel brokerageModel, ISecuritySeeder securitySeeder)
                : base(brokerageModel, securitySeeder) {}    

            public override void Initialize(Security security)
            {
                base.Initialize(security);
                // Set the shortable provider as your broker for accurate short reality modeling
                security.SetShortableProvider(new InteractiveBrokersShortableProvider());    
            }
        }
    }
}

from AlgorithmImports import *
from QuantConnect.DataSource import *

class ShortAvailabilityDataAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        
        self.set_start_date(2021, 1, 1)
        self.set_end_date(2021, 7, 1)
        self.set_cash(1000)
        # Seed the security price as the last known price, such that the price data is immediately available at initial rebalance
        self.set_security_initializer(MySecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        self.set_universe_selection(ManualUniverseSelectionModel(
            [Symbol.create("GME", SecurityType.EQUITY, Market.USA)]))           

        # Emit down-direction insights to short all securities in the universe            
        self.set_alpha(ConstantAlphaModel(InsightType.PRICE, InsightDirection.DOWN, timedelta(1)))
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel())
        self.set_execution(ShortableExecutionModel())

    # On Margin Call, emit flat-direction insights to liquidate the positions
    def on_margin_call_warning(self) -> None:
        self.emit_insights([Insight.price(kvp.Key, timedelta(1), InsightDirection.FLAT)
            for kvp in self.securities if kvp.Value.invested])

class ShortableExecutionModel(ExecutionModel):
    def __init__(self) -> None:
        self.targets_collection = PortfolioTargetCollection()

    def execute(self, algorithm: QCAlgorithm, targets: List[PortfolioTarget]) -> None:
        '''Immediately submits orders for the specified portfolio targets.
        Args:
            algorithm: The algorithm instance
            targets: The portfolio targets to be ordered'''

        # for performance we check count value, OrderByMarginImpact and ClearFulfilled are expensive to call
        self.targets_collection.add_range(targets)
        if self.targets_collection.count > 0:
            for target in self.targets_collection.order_by_margin_impact(algorithm):
                # calculate remaining quantity to be ordered
                quantity = OrderSizing.get_unordered_quantity(algorithm, target)
                # If the quantity is negative, ensure that the security is shortable
                if quantity > 0 or algorithm.shortable(target.symbol, quantity):
                    algorithm.market_order(target.symbol, quantity)

            self.targets_collection.clear_fulfilled(algorithm)

class MySecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(self, brokerage_model: IBrokerageModel, security_seeder: ISecuritySeeder) -> None:
        super().__init__(brokerage_model, security_seeder)

    def initialize(self, security: Security) -> None:
        super().initialize(security)
        # Set the shortable provider as IB
        security.set_shortable_provider(InteractiveBrokersShortableProvider())

using QuantConnect.Data.Shortable;
namespace QuantConnect
{
    public class ShortAvailabilityDataAlgorithm : QCAlgorithm
    {
        public override void Initialize()
        {
            SetStartDate(2021, 1, 1);
            SetEndDate(2021, 7, 1);
            SetCash(1000);
            // Seed the security price as the last known price, such that the price data is immediately available at initial rebalance
            SetSecurityInitializer(new MySecurityInitializer(BrokerageModel, new FuncSecuritySeeder(GetLastKnownPrices)));

            SetUniverseSelection(new ManualUniverseSelectionModel(
                QuantConnect.Symbol.Create("GME", SecurityType.Equity, Market.USA)));            
            
            // Emit down-direction insights to short all securities in the universe            
            SetAlpha(new ConstantAlphaModel(InsightType.Price, InsightDirection.Down, TimeSpan.FromDays(1)));
            SetPortfolioConstruction(new EqualWeightingPortfolioConstructionModel());
            SetExecution(new ShortableExecutionModel());
        }

        // On Margin Call, emit flat-direction insights to liquidate the positions
        public override void OnMarginCallWarning()
        {
            EmitInsights(Securities
                .Where(kvp => kvp.Value.Invested)
                .Select(kvp => Insight.Price(kvp.Key, TimeSpan.FromDays(1), InsightDirection.Flat))
                .ToArray());
        }

        public class ShortableExecutionModel : ImmediateExecutionModel
        {
            private readonly PortfolioTargetCollection _targetsCollection = new PortfolioTargetCollection();

            public override void Execute(QCAlgorithm algorithm, IPortfolioTarget[] targets)
            {
                _targetsCollection.AddRange(targets);
                // for performance we check count value, OrderByMarginImpact and ClearFulfilled are expensive to call
                if (_targetsCollection.Count > 0)
                {
                    foreach (var target in _targetsCollection.OrderByMarginImpact(algorithm))
                    {
                        // calculate remaining quantity to be ordered
                        var quantity = OrderSizing.GetUnorderedQuantity(algorithm, target);
                        // If the quantity is negative, ensure that the security is shortable
                        if (quantity > 0 || algorithm.Shortable(target.Symbol, quantity))
                        {
                            algorithm.MarketOrder(target.Symbol, quantity);
                        }
                    }

                    _targetsCollection.ClearFulfilled(algorithm);
                }
            }
        }

        class MySecurityInitializer : BrokerageModelSecurityInitializer
        {
            public MySecurityInitializer(IBrokerageModel brokerageModel, ISecuritySeeder securitySeeder)
                : base(brokerageModel, securitySeeder) {}    

            public override void Initialize(Security security)
            {
                base.Initialize(security);
                // Set the shortable provider as IB
                security.SetShortableProvider(new InteractiveBrokersShortableProvider());    
            }
        }
    }
}