from AlgorithmImports import *
from QuantConnect.DataSource import *


class EstimizeDeployedApiProbe(QCAlgorithm):
    """Runs the demo's logic against the library the cloud has deployed today, where
    EstimizeConsensus/EstimizeEstimate are themselves the collections (d7dbbe3)."""

    def initialize(self):
        self.set_start_date(2019, 1, 1)
        self.set_end_date(2019, 12, 31)
        self.set_cash(100_000)

        self._equity = self.add_equity("AAPL", Resolution.DAILY).symbol
        self._consensus = self.add_data(EstimizeConsensus, self._equity).symbol
        self._estimate = self.add_data(EstimizeEstimate, self._equity).symbol

        self._points = 0
        self._instants = 0
        self._reported = False

    def on_data(self, slice):
        if slice.contains_key(self._consensus):
            collection = slice[self._consensus]
            self._instants += 1
            self._points += len(collection.data)

            if not self._reported and len(collection.data) > 0:
                self._reported = True
                point = collection.data[0]
                self.log("POINT attrs: " + ", ".join(sorted(
                    a for a in dir(point) if not a.startswith("_") and a.islower())))

            latest_release = max(p.id for p in collection.data)
            estimize_eps = None
            wall_street_eps = None
            for p in collection.data:
                if p.id != latest_release or p.type != EstimizeConsensus.ConsensusType.EPS:
                    continue
                if p.source == EstimizeConsensus.ConsensusSource.ESTIMIZE:
                    estimize_eps = p.mean
                elif p.source == EstimizeConsensus.ConsensusSource.WALL_STREET:
                    wall_street_eps = p.mean

            if self._instants <= 5:
                self.log(f"{self.time} {len(collection.data)} points, release {latest_release}, "
                         f"estimize {estimize_eps}, street {wall_street_eps}")

            if estimize_eps is not None and wall_street_eps is not None:
                if estimize_eps > wall_street_eps:
                    self.set_holdings(self._equity, 1)
                else:
                    self.liquidate(self._equity)

        if slice.contains_key(self._estimate):
            collection = slice[self._estimate]
            eps = [p.eps for p in collection.data if p.eps is not None]
            if len(eps) >= 2 and self._instants <= 5:
                self.log(f"{self.time} {len(collection.data)} analysts, EPS {min(eps)} to {max(eps)}")

    def on_end_of_algorithm(self):
        self.log(f"instants: {self._instants}, consensus points: {self._points}")