# region imports
from AlgorithmImports import *
# endregion


class NorthFieldAlpha(PythonData):

    def get_source(self, config: SubscriptionDataConfig, date: datetime, is_live: bool):
        return SubscriptionDataSource(
            f"Alphas/Alphas {date.strftime('%Y%m%d')}.txt",
            SubscriptionTransportMedium.OBJECT_STORE,
            FileFormat.CSV
        )

    def reader(self, config: SubscriptionDataConfig, line: str, date: datetime, is_live: bool):
        if not line[0].isdigit():
            return None
        data = line.split('|')
        asset = NorthFieldAlpha()
        asset.symbol = SecurityDefinitionSymbolResolver.get_instance().cusip(data[0], date)
        if not asset.symbol:
            return
        asset.name = data[1]
        try:
            asset.alpha = float(data[2]) # These fields can be empty in the data file.
            asset.market_cap = float(data[3])
        except:
            return None
        asset.end_time = date #+ timedelta(1)
        return asset

# region imports
from AlgorithmImports import *
# endregion

class NorthFieldFactorCorrelation(PythonData):

    _rows_seen_by_date = {}

    def get_source(self, config: SubscriptionDataConfig, date: datetime, is_live: bool):
        return SubscriptionDataSource(
            f"US_2_19_9g/FF_RSQ_RSQRM_US_v2_19_9g_USD_{date.strftime('%Y%m%d')}_Correl.txt",
            SubscriptionTransportMedium.OBJECT_STORE,
            FileFormat.CSV
        )

    def reader(self, config: SubscriptionDataConfig, line: str, date: datetime, is_live: bool):
        data = line.split('|')
        try:
            float(data[0])
        except:
            self._columns = data
            return None

        if date not in self._rows_seen_by_date:
            self._rows_seen_by_date[date] = 0
        else:
            self._rows_seen_by_date[date] += 1

        fc = NorthFieldFactorCorrelation()
        factor_name = self._columns[self._rows_seen_by_date[date]].replace(' ', '')
        fc.symbol = Symbol.create(factor_name, SecurityType.BASE, Market.USA)
        fc.end_time = date
        
        # Parse columns.
        fc.series = pd.Series(data, index=self._columns).astype(float)

        return fc

# region imports
from AlgorithmImports import *
# endregion

class NorthFieldFactorDefinition(PythonData):

    def get_source(self, config: SubscriptionDataConfig, date: datetime, is_live: bool):
        return SubscriptionDataSource(
            f"US_2_19_9g/FF_RSQ_RSQRM_US_v2_19_9g_USD_{date.strftime('%Y%m%d')}_FactorDef.txt",
            SubscriptionTransportMedium.OBJECT_STORE,
            FileFormat.CSV
        )

    def reader(self, config: SubscriptionDataConfig, line: str, date: datetime, is_live: bool):
        if not line[0].isdigit():
            return None
        data = line.split('|')
        fd = NorthFieldFactorDefinition()
        
        # Parse columns.
        fd.name = data[1]
        fd.code = data[2]
        fd.variance = float(data[3])

        fd.symbol = Symbol.create(fd.code, SecurityType.BASE, Market.USA, fd.name)
        fd.end_time = date

        return fd

# region imports
from AlgorithmImports import *

import csv
from io import StringIO
# endregion

class NorthFieldFactorExposure(PythonData):

    _drop_unsupported_assets = False  # Change this if you want.

    def get_source(self, config: SubscriptionDataConfig, date: datetime, is_live: bool):
        return SubscriptionDataSource(
            f"US_2_19_9g/RSQRM_US_v2_19_9g{date.strftime('%Y%m%d')}_USDcusip.csv",
            SubscriptionTransportMedium.OBJECT_STORE,
            FileFormat.CSV
        )

    def reader(self, config: SubscriptionDataConfig, line: str, date: datetime, is_live: bool):       
        data = next(csv.reader(StringIO(line)))
        fe = NorthFieldFactorExposure()
        
        # Parse columns.
        fe.name = data[1]
        fe.value = data[3]
        fe.monthly_residual_sd_pct = float(data[4])
        fe.betas = [float(data[i]) for i in range(6, 47)]
        fe.currency_quotation = data[48]
        fe.total_forecast_risk_ann_pct = data[49]

        # QC doesn't have ADRs, OTC stocks, European stocks, Forex without quote currencies, industry proxies...
        if "*" in data[0] or "_" in data[0]:
            fe.symbol = Symbol.create(data[0], SecurityType.BASE, Market.USA)
        else:
            fe.symbol = SecurityDefinitionSymbolResolver.get_instance().cusip(data[0], date)

        if not fe.symbol:
            #NorthFieldFactorExposure.algorithm.log(f"Error creating Symbol for {data[0]} at {date}. Name: {fe.name}")
            # Example CUSIPs that are dropped:
            # 2024-06-05 00:00:00 Error creating Symbol for G9T17Y80 at 2024-06-05 00:00:00. Name: "VANGD.US TRSY.0-1Y (OTC) BD UCITS ETF USD ACC"
            # 2024-06-05 00:00:00 Error creating Symbol for 06254520 at 2024-06-05 00:00:00. Name: "BANK OF HAWAII DRC EACH"
            # 2024-06-05 00:00:00 Error creating Symbol for 57142B10 at 2024-06-05 00:00:00. Name: "MARQETA A"
            # 2024-06-05 00:00:00 Error creating Symbol for G6543112 at 2024-06-05 00:00:00. Name: "NOBLE CORPORATION"
            # 2024-06-05 00:00:00 Error creating Symbol for 87283Q50 at 2024-06-05 00:00:00. Name: "T ROWE PRICE US EQUITY RESEARCH ETF"
            # 2024-06-05 00:00:00 Error creating Symbol for 50076757 at 2024-06-05 00:00:00. Name: "KRANESHARES HANG SENG TECH INDEX ETF"
            # 2024-06-05 00:00:00 Error creating Symbol for 29482Y20 at 2024-06-05 00:00:00. Name: "ERICKSON"
            # 2024-06-05 00:00:00 Error creating Symbol for 44615078 at 2024-06-05 00:00:00. Name: "HUNTINGTON BANCSHARES DEP"
            # 2024-06-05 00:00:00 Error creating Symbol for 41150T20 at 2024-06-05 00:00:00. Name: "HARBOR CUSTOM DEV.8 0 CUM CONV PREF. SR.A"
            if self._drop_unsupported_assets:
                return None
            else:
                fe.symbol = Symbol.create(data[0], SecurityType.BASE, Market.USA)

        fe.end_time = date
        
        return fe

# region imports
from AlgorithmImports import *
from alpha import NorthFieldAlpha
# endregion

# Documentation:
# - README: https://www.dropbox.com/scl/fi/5lubbi4p7b9art7ig73v6/read-me.docx?rlkey=bwvgomxo4b1kh3dtkzf8mvdtm&st=bz0ja53n&dl=0
# - Flat File description: https://www.dropbox.com/scl/fi/s5mgdu69lw7ef5k6wfs00/Flat-File-description-FF_-prefix.pdf?rlkey=uywqpaw8wvpw26no3isomqn5y&st=h0waii0g&dl=0

# Demo 4: Top 10 based on Alpha
class NorthFieldDemoAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2017, 12, 20)
        self.set_end_date(2024, 8, 15)
        self.universe_settings.resolution = Resolution.MINUTE

        # Add alpha universe.
        self._universe = self.add_universe(NorthFieldAlpha, 'NorthFieldAlpha', Resolution.DAILY, self._select_assets)

        spy = Symbol.create("SPY", SecurityType.EQUITY, Market.USA)
        self.schedule.on(self.date_rules.every_day(spy), self.time_rules.after_market_open(spy, 30), self.rebalance)

    def _select_assets(self, data):
        data = sorted(data, key=lambda x: x.alpha, reverse=True)[:10]
        return [x.symbol for x in data]

    def on_securities_changed(self, changes):
        for security in changes.removed_securities:
            self.liquidate(security.symbol)

    def rebalance(self):
        weight = 1 / len(self.active_securities.keys) if self.active_securities.keys else 0
        self.set_holdings([PortfolioTarget(symbol, weight) for symbol in self.active_securities.keys])

# region imports
from AlgorithmImports import *
# endregion


class NorthFieldUniverse(PythonData):

    def get_source(self, config: SubscriptionDataConfig, date: datetime, is_live: bool):
        return SubscriptionDataSource(
            f"GeneralUniverse/Universe {date.strftime('%Y%m%d')}.txt",
            SubscriptionTransportMedium.OBJECT_STORE,
            FileFormat.CSV
        )

    def reader(self, config: SubscriptionDataConfig, line: str, date: datetime, is_live: bool):
        if not line[0].isdigit():
            return None
        data = line.split('|')
        asset = NorthFieldUniverse()
        symbol = SecurityDefinitionSymbolResolver.get_instance().cusip(data[0], date)
        if not symbol:
            return
        asset.symbol = symbol
        asset.market_cap = float(data[1])
        asset.name = data[2]
        asset.value = asset.market_cap
        asset.end_time = date
        return asset


class NorthFieldInvestableUniverse(NorthFieldUniverse):

    def get_source(self, config: SubscriptionDataConfig, date: datetime, is_live: bool):
        return SubscriptionDataSource(
            f"InvestableUniverse/InvestableUniverse {date.strftime('%Y%m%d')}.txt",
            SubscriptionTransportMedium.OBJECT_STORE,
            FileFormat.CSV
        )