# region imports
from AlgorithmImports import *
# endregion

class UpgradedBrownElephant(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2008, 12, 31)
        self.set_end_date(2023, 1, 1) # The dataset hasn't been updated after this point.
        self._tariffs = self.add_data(Tariffs, "Tariffs", Resolution.DAILY).symbol
        self.set_cash(1_000_000)

        etf_by_country = {
            "Croatia" : "FM",
            "Singapore" : "EWS",
            "Hungary" : "CRAK",
            "Denmark" : "EDEN",
            "Norway" : "NORW",
            "United Kingdom" : "EWUS",
            "Sweden" : "EWD",
            "France" : "FLFR",
            "Luxembourg" : "SLX",
            "Japan" : "DFJ",
            "Germany" : "FLGR",
            "Greece" : "GREK",
            "Finland" : "EFNL",
            "Ireland" : "EIRL",
            "UAE" : "UAE",
            "Qatar" : "QAT",
            "Czech Republic" : "NLR",
            "Portugal" : "PGAL",
            "Austria" : "EWO",
            "South Korea" : "FLKR",
            "Hong Kong" : "EWH",
            "Belgium" : "EWK",
            "Cyprus" : "SIL",
            "Switzerland" : "EWL",
            "United States": "SPY",
            "Malta" : "BETZ",
            "Canada" : "FLCA",
            "Israel" : "EIS",
            "Spain" : "EWP",
            "Poland" : "EPOL",
            "Argentina" : "ARGT",
            "Estonia" : "FM",
            "Netherlands" : "EWN",
            "Australia" : "EWA",
            "Italy" : "FLIY",
            'Trinidad and Tobago': 'FDIV',
            'Chile': 'ECH',
            'New Zealand': 'ENZL',
            'Oman': 'FDIV', 
            'Uruguay': 'FDIV',
            'Macao': 'MOTI',
            'Panama': 'TUR',
            'Bahrain': 'KWT',
            'Kuwait': 'KWT',
            'Bermuda': 'EQIN',
            'Cayman Islands': 'EMM'
        }
        self._symbol_by_country = {
            country: self.add_equity(ticker).symbol
            for country, ticker in etf_by_country.items()
        }

    def on_data(self, data):
        # Wait until there is new tariff data released.
        if self._tariffs not in data:
            return
        tariffs = data[self._tariffs]
        # Select the subset of countries where their ETF has a price already.
        tradable = []
        for x in tariffs['country_data']:
            if x['country'] not in self._symbol_by_country:
                continue
            if not self.securities[self._symbol_by_country[x['country']]].price:
                continue
            tradable.append(x)
        if not tradable:
            return
        # Rebalance the portfolio.
        tariffs_sum = sum([x['value'] for x in tradable])
        weight_by_symbol = {}
        for x in tradable:
            symbol = self._symbol_by_country[x['country']]
            if symbol not in weight_by_symbol:
                weight_by_symbol[symbol] = 0
            weight_by_symbol[symbol] += (x['value'] / tariffs_sum)
        self.set_holdings([PortfolioTarget(symbol, weight) for symbol, weight in weight_by_symbol.items()], True)

class Tariffs(PythonData):

    def get_source(self, config: SubscriptionDataConfig, date: datetime, is_live: bool) -> SubscriptionDataSource:
        return SubscriptionDataSource("tariff_data.json", SubscriptionTransportMedium.OBJECT_STORE, FileFormat.UNFOLDING_COLLECTION)

    def reader(self, config: SubscriptionDataConfig, line: str, date: datetime, is_live: bool) -> BaseData:
        objects = []
        samples = json.loads(line)
        for sample in samples:
            tariffs = Tariffs()
            tariffs.symbol = config.symbol
            tariffs['year'] = sample['year']
            tariffs['country_data'] = [
                {
                    'country': country_info['country'],
                    'value': country_info['tariff'],
                    'annual_change': country_info['annual_change'] # 0.01 = 1%
                }
                for country_info in sample['data']
            ]
            tariffs.value = 0
            tariffs.end_time = datetime(tariffs['year']+1, 1, 1)
            objects.append(tariffs)

        return BaseDataCollection(objects[-1].end_time, config.symbol, objects)