# region imports
from AlgorithmImports import *
# endregion

class FocusedRedLeopard(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2010, 1, 1)

        self.SetCash(100000)

        self.future = self.AddFuture(
            Futures.Financials.Y10TreasuryNote, 
            dataMappingMode=DataMappingMode.LastTradingDay, 
            dataNormalizationMode=DataNormalizationMode.BackwardsPanamaCanal
        )
        self.future.SetFilter(0, 400)
        self.day = 0

        self.annualized_raw_carry = pd.Series()
        self.raw_carry = pd.Series()
        self.raw_price = pd.Series()
        self.adjusted_price = pd.Series()

    def OnData(self, data: Slice):
        if self.day == self.Time.day:
            return
        self.day = self.Time.day

        near_contract, further_contract = self.get_near_and_further_contracts(self.Securities, self.future.Mapped)
        raw_carry = near_contract.Price - further_contract.Price
        self.raw_carry.loc[self.Time] = raw_carry
        self.raw_price.loc[self.Time] = self.Securities[self.future.Mapped].Price
        self.adjusted_price.loc[self.Time] = self.Securities[self.future.Symbol].Price

        # Annualize raw carry
        months_between_contracts = round((further_contract.Expiry - near_contract.Expiry).days / 30)
        expiry_difference_in_years = abs(months_between_contracts) / 12
        annualized_raw_carry = raw_carry / expiry_difference_in_years
        self.annualized_raw_carry.loc[self.Time] = annualized_raw_carry

    def get_near_and_further_contracts(self, securities, mapped_symbol):
        ## Gather and align history of near/further contracts
        contracts_sorted_by_expiry = sorted(
            [
                kvp.Value for kvp in securities 
                if not kvp.Key.IsCanonical() \
                    and kvp.Key.Canonical == mapped_symbol.Canonical \
                    and kvp.Value.Expiry >= securities[mapped_symbol].Expiry
            ], 
            key=lambda contract: contract.Expiry
        )
        near_contract = contracts_sorted_by_expiry[0]
        further_contract = contracts_sorted_by_expiry[1]
        return near_contract, further_contract

    def OnEndOfAlgorithm(self):
        self.ObjectStore.Save(f"{self.ProjectId}/raw_carry", self.raw_carry.to_json())
        self.ObjectStore.Save(f"{self.ProjectId}/annualized_raw_carry", self.annualized_raw_carry.to_json())
        self.ObjectStore.Save(f"{self.ProjectId}/raw_price", self.raw_price.to_json())
        self.ObjectStore.Save(f"{self.ProjectId}/adjusted_price", self.adjusted_price.to_json())

# region imports
from AlgorithmImports import *

from Selection.FutureUniverseSelectionModel import FutureUniverseSelectionModel
from futures import categories
# endregion

class AdvancedFuturesUniverseSelectionModel(FutureUniverseSelectionModel):
    def __init__(self) -> None:
        super().__init__(timedelta(1), self.select_future_chain_symbols)
        self.symbols = list(categories.keys())

    def select_future_chain_symbols(self, utc_time: datetime) -> List[Symbol]:
        return self.symbols

    def Filter(self, filter: FutureFilterUniverse) -> FutureFilterUniverse:
        return filter.Expiration(0, 200)

# region imports
from AlgorithmImports import *
# endregion


class FutureData:
    def __init__(self, classification, contract_offset):
        self.classification = classification
        self.contract_offset = contract_offset

categories = {
    pair[0]: FutureData(pair[1], pair[2]) for pair in [
        #(Symbol.Create(Futures.Energies.NaturalGas, SecurityType.Future, Market.NYMEX), ("Energies", "Gas"), 1),
        (Symbol.Create(Futures.Indices.SP500EMini, SecurityType.Future, Market.CME), ("Equity", "US"), 0),
    ]
}