from typing import List, Dict
from datetime import timedelta
import pandas as pd
import numpy as np

from QuantConnect import Resolution
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data import Slice
from QuantConnect.Data.Consolidators import TradeBarConsolidator, DataConsolidator
from QuantConnect.Securities import Futures
from QuantConnect.Securities.Future import Future
from QuantConnect.Data.Market import FuturesContract, FuturesChain, FuturesChains, FuturesContracts, TradeBar


class ParticleTachyonInterceptor(QCAlgorithm):
    consolidator_by_symbol = {}
    contract_to_parent_map = {}
    parent_to_active_map = {}
    contracts_info = {}
    universe: Dict[str, Future]
    universe = {}

    db = pd.DataFrame(columns=["parent_symbol", "contract_symbol", "close", "oi", "end_time"])
    continuous_futures = pd.DataFrame()

    def Initialize(self):
        self.SetStartDate(2020, 5, 31)
        self.SetEndDate(2020, 10, 31)

        self.SetCash(1000000)  # Set Strategy Cash
        self.universe = {symbol: self.AddFuture(symbol, Resolution.Minute) for symbol in self.tickers}
        [future.SetFilter(timedelta(0), timedelta(90)) for future in self.universe.values()]

        self.consolidator_by_symbol = {}
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(6, 0), self.morning_routine)

    def morning_routine(self):
        self.update_continuous_prices()
        self.Debug(f"Status update on {self.Time}")
        if not self.continuous_futures.empty:
            self.Debug(self.continuous_futures.describe().to_string())

    def update_continuous_prices(self):
        unique_futures = self.db.parent_symbol.unique()
        continuous_data = []
        if self.db.empty:
            return
        for symbol in unique_futures:
            symbol_data = self.db.loc[self.db.parent_symbol == symbol]
            oi = symbol_data.pivot_table(index="end_time", columns="contract_symbol", values="oi")
            price = symbol_data.pivot_table(index="end_time", columns="contract_symbol", values="close")
            weights = oi.divide(oi.sum(axis=1), axis=0)
            combined = (price * weights).sum(axis=1).rename(symbol.Value)
            continuous_data.append(combined)
        if len(continuous_data) > 0:
            self.continuous_futures = pd.concat(continuous_data, axis=1, sort=False)

    def OnData(self, data: Slice):
        chain: FuturesChain
        contracts: FuturesContracts
        contract: FuturesContract
        sorted_by_oi_contracts: List[FuturesContract]

        for chain in data.FutureChains.Values:
            contracts = chain.Contracts
            if len(contracts) == 0:
                continue
            sorted_by_oi_contracts = sorted(contracts.Values, key=lambda k: k.OpenInterest, reverse=True)

            popular_contracts = sorted_by_oi_contracts[:3]
            active_contract = sorted_by_oi_contracts[0]
            # Record active contract
            self.parent_to_active_map[chain.Symbol] = active_contract.Symbol
            # Record popular contracts map
            for contract in popular_contracts:
                if contract.Symbol not in self.contract_to_parent_map:
                    self.contract_to_parent_map[contract.Symbol] = chain.Symbol
                if contract.Symbol not in self.consolidator_by_symbol:
                    self.consolidate_future(contract)
                self.contracts_info[contract.Symbol] = contract

    def consolidate_future(self, contract):
        dailyConsolidator = TradeBarConsolidator(timedelta(days=1))
        dailyConsolidator.DataConsolidated += self.DailyHandler
        self.SubscriptionManager.AddConsolidator(contract.Symbol, dailyConsolidator)
        self.consolidator_by_symbol[contract.Symbol] = dailyConsolidator

    def DailyHandler(self, sender: DataConsolidator, bar: TradeBar):
        contract_symbol = sender.Consolidated.Symbol
        parent_symbol = self.contract_to_parent_map[contract_symbol]
        end_time = bar.EndTime
        close = bar.Close
        oi = self.contracts_info[contract_symbol].OpenInterest
        data = pd.Series({
            "parent_symbol": parent_symbol,
            "contract_symbol": contract_symbol,
            "close": close,
            "oi": oi,
            "end_time": end_time
        })
        self.db = self.db.append(data, ignore_index=True)

    @property
    def tickers(self):
        indices = [
            Futures.Indices.SP500EMini,
            Futures.Indices.SP400MidCapEmini,
            Futures.Indices.Dow30EMini,
            Futures.Indices.Russell2000EMini,

            Futures.Indices.Nikkei225Dollar,
            Futures.Indices.USDDenominatedIbovespa,

            Futures.Indices.SPGSCICommodity,
            Futures.Indices.BloombergCommodityIndex,

            Futures.Indices.VIX,
        ]

        metals = [
            Futures.Metals.Gold,
            Futures.Metals.Silver,
            Futures.Metals.Platinum,
            Futures.Metals.Palladium,
            Futures.Metals.Copper,
        ]

        grains = [
            Futures.Grains.Wheat,
            Futures.Grains.Corn,
            Futures.Grains.Soybeans,
            Futures.Grains.Oats,
        ]

        tickers = indices + metals + grains

        return tickers