from AlgorithmImports import *
from prediction_data import PredictionData
from test import add
import pandas as pd

def print_file(algorithm):
    path = algorithm.object_store.get_file_path("processed_monthly_2024-12-05.csv")
    df = pd.read_csv(path)
    algorithm.quit(df.head().to_string())

class SectorBasedTradingAlgorithm(QCAlgorithm):
    def initialize(self):
        self.name = f'MyBacktest_top10'
        #print_file(self)
        self.set_start_date(2024, 12, 1)
        self.set_end_date(2024, 12, 31)
        self.set_cash(50_000_000)
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        self.add_universe(PredictionData, "PredictionData", Resolution.DAILY, self._selector_function)
        self.content = ''
        self.trades = []

    def _selector_function(self, data: List[PredictionData]) -> List[Symbol]:
        selected = sorted(data, key=lambda x: x.pred_proba)[:10]
        self.trades = selected
        return [x.symbol for x in self.trades]

    def on_data(self, data):

        # We will iterate trades until all symbols are invested
        for trade in self.trades:
            symbol = trade.symbol
            if self.portfolio[symbol].invested:
                continue
            bar = data.bars.get(symbol)
            if not bar:
                continue
            if trade.side == OrderDirection.SELL:
                self.sell(symbol, trade.quantity)
            if trade.side == OrderDirection.BUY:
                self.buy(symbol, trade.quantity)

    def on_end_of_day(self, symbol):
        holdings = self.portfolio[symbol]
        self.content += f'{self.time},{holdings.symbol},{holdings.profit}'

    def on_end_of_algorithm(self):
        self.object_store.save('stats', self.content)

    def on_securities_changed(self, changes):
        for security in changes.removed_securities:
            self.liquidate(security.symbol, tag='Removed from Universe')

# region imports
from AlgorithmImports import *
# endregion

class PredictionData(PythonData):
    def get_source(self, config, date, isLiveMode):
        key = f'processed_monthly_{date:%Y-%m-%d}.csv'
        #key = 'processed_monthly_2024-12-05.csv'
        return SubscriptionDataSource(key, SubscriptionTransportMedium.OBJECT_STORE)

    def reader(self, config, line, date, isLiveMode):
        # Example Line Format:
        # id,model_id,ticker,prediction_date,close,pred_proba,prediction,side,Quantity
        # 797131,319,NSP,12/4/24,82.87999725,0.869273663,-1,Sell,600
        if not line[0].isdigit():
            return None

        csv = line.split(',')
        quantity = float(csv[8])
        if quantity == 0:
            return None
        time = date #datetime.strptime(csv[3], "%m/%d/%y")
        sid = SecurityIdentifier.generate_equity(csv[2], Market.USA, mapping_resolve_date=time)

        data = PredictionData()
        data.symbol = Symbol(sid, csv[2])
        data.time = time - timedelta(1)
        data.end_time = time
        data.value = float(csv[4])

        data['id'] = csv[0]
        data['model_id'] = csv[1]
        data['pred_proba'] = float(csv[5])
        data['prediction'] = float(csv[6])
        data['quantity'] = quantity

        match csv[7].lower().strip():
            case 'sell':
                data['side'] = OrderDirection.SELL
            case 'buy':
                data['side'] = OrderDirection.BUY
            case _:
                data['side'] = OrderDirection.HOLD

        return data

# region imports
from AlgorithmImports import *
# endregion

# Your New Python File
def add(a, b):
    return a+b