import pandas as pd
import numpy as np
import math
from datetime import datetime, timedelta

### INITIALIZE

# Variables
stocks = ["SPK","WPP","CCV","DVN","CYG","SDI","DTL","MAH","REH","TAH","GRR","AIA","NCM","ORI","AVJ","MCR","CLV","IGO","GME","ALL","HAV","SEN","WEB","MGX","WOR","SBM","OZL","IMD","AUZ","ANG","RMS","PAN","OMH","CVN","NHC","MTS","X64","JBH","SLR","QTM","OGC","CCX","FLT","GAP","MAY","VMG","AMA","DDT","MLX","GEM","RRL","ILU","VYS","CGM","PHX","JIN","SYD","AGI","CAJ","IFM","NEA","IMS","VMT","GOR","RHT","ELD","VPR","EGH","ADA","BAR","NWH","NAE","BSE","WHC","BCN","LYC","OLI","BAL","BYE","EHL","DCC","ACF","RED","ALK","M7T","HTG","TPW","HT8","ZNO","MGV","ICI","RNT","MAI","FSG","AIS","BRK","ATP","TER","SMR","NEU","CTT","SMP","MIL","MP1","CBA","DLC","ASB","SGP","SWM","TNE","IRI","RDF","AAR","IPL","CIM","AJL","APA","ACR","NST","FDM","PDZ","CDA","FNP","CSS","RHP","QAN","PXS","CMP","SSM","LPD","AMI","PYC","WTC","CGR","ADH","LNY","PME","ZIP","CXZ","AEF","ESK","OEC","TNT","JPR","AQS","LCY","NGY","EGL","IDT","JLG","CRN","BRL","YAL","DUR","LAU","DUG","360","SVW","CIA","BLD","A2B","JHX",]
portfolio_value = 10000                         
weight = 0.10                                   
start_date = "2002-01-31"                       
end_date = "2023-12-31"    

# Universe selection
universe = pd.read_csv("C:/Users/jerem/LeanCLI/data/asx_fundamentals\ASX Piotroski Universe Selection.csv", header=None, index_col=0)
universe.fillna('', inplace=True)


# Gather price data
stock_data = pd.DataFrame(columns=stocks)
stock_data['date'] = []
stock_data.set_index('date', inplace=True)
for stock in stocks:
    data = pd.read_csv("C:/Users/jerem/LeanCLI/data/asx_prices/{} Price Data.csv".format(stock), index_col='date')
    stock_data[stock] = data['adjusted_close']

# Set initial shares on first day
shares_df = pd.DataFrame(index=[universe.loc[start_date].index])
for s in universe.loc[start_date]:
    if s != '':
        shares_df[s + '_shares'] = np.floor((portfolio_value * weight) / stock_data[s][0])

### REBALANCING ENGINE

# Variables
balance_date = datetime.strptime(start_date, "%Y-%m-%d")
balance_month = balance_date.month
count = 0
prev_values = {}

"""for day in stock_data.index:
    count += 1
    if day.month != balance_month:
        signal = True
        new_shares = universe.loc[balance_date]"""

print(shares_df)

# region imports
from datetime import date, datetime, timedelta
from AlgorithmImports import *
from QuantConnect.Data import BaseData, SubscriptionDataConfig
from QuantConnect.Python import PythonData
from QuantConnect import Globals, SubscriptionTransportMedium
import os
# endregion

class ASXPiotroski(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2002, 1, 1)
        self.SetEndDate(2023, 12, 31)
        self.SetAccountCurrency("AUD")
        self.SetCash(10000)
        self.SetBrokerageModel(BrokerageName.QuantConnectBrokerage, AccountType.Margin)
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        
        # Add Custom Universe
        self.AddUniverse(StockDataSource, "ASX_Piotroski_Universe", self.selector_function)
        self.UniverseSettings.Resolution = Resolution.Daily
        
        # Variables
        self.stoploss = -0.1
        self.num_fine = 10
        self.data_uploaded = []

        # Rebalancing
        self.month = 0

    def CustomSecurityInitializer(self, security):
        security.SetLeverage(2)

    def selector_function(self, data):

        list = []
        for item in data:
            for symbol in item["Symbols"]:
                list.append(symbol)
                if symbol not in self.data_uploaded:
                    self.AddData(CustomOHLC, symbol, Resolution.Daily)
                    self.data_uploaded.append(symbol)
        return list

    def OnData(self, slice: Slice):
        
        if self.IsWarmingUp: return

        for security in self.ActiveSecurities.Values:
            if slice.ContainsKey(security.Symbol):
                custom_data = slice[security.Symbol]
                close = custom_data.Value
            else:
                return

        for symbol in self.Portfolio.Keys:
                pnl = self.Securities[symbol].Holdings.UnrealizedProfitPercent
                if pnl < self.stoploss:
                    self.Liquidate(symbol, "Stop Loss initiated.")
    
    def OnSecuritiesChanged(self, changes):
                    
        # Liquidate removed securities
        for security in changes.RemovedSecurities:
            self.Liquidate(security.Symbol, "Removed from universe")
            self.Log(str(security.Symbol) + " has been removed from the universe.")

        # Purchase added securities
        for security in changes.AddedSecurities:
            self.SetHoldings(security.Symbol, 1/self.num_fine * security.Leverage, False, "Added to universe")
            self.Log(str(security.Symbol) + " has been added to the universe.")

class StockDataSource(PythonData):

    def GetSource(self,
         config: SubscriptionDataConfig,
         date: datetime,
         isLiveMode: bool) -> SubscriptionDataSource:
        
        source = os.path.join(Globals.DataFolder,"asx_fundamentals/ASX Piotroski Universe Selection.csv")
        return SubscriptionDataSource(source, SubscriptionTransportMedium.LocalFile)
    
    def Reader(self,
         config: SubscriptionDataConfig,
         line: str,
         date: datetime,
         isLive: bool) -> BaseData:
        
        if not (line.strip() and line[0].isdigit()): return None

        stocks = StockDataSource()
        stocks.Symbol = config.Symbol

        try:
            csv = line.split(',')
            csv = list(filter(None,csv))
            stocks.Time = datetime.strptime(csv[0], "%Y-%m-%d")
            stocks.EndTime = stocks.Time + timedelta(days=1)
            stocks["Symbols"] = csv[1:]
        
        except ValueError:
            return None
        
        return stocks
    
class CustomOHLC(PythonData):
    
    def GetSource(self,
         config: SubscriptionDataConfig,
         date: datetime,
         isLiveMode: bool) -> SubscriptionDataSource:
        
        symbol = config.Symbol.Value.split(" ")[0]
        source = os.path.join(Globals.DataFolder,"asx_prices/{} Price Data.csv".format(symbol))
        return SubscriptionDataSource(source, SubscriptionTransportMedium.LocalFile)
    
    def Reader(self,
         config: SubscriptionDataConfig,
         line: str,
         date: datetime,
         isLive: bool) -> BaseData:
        
        if not (line.strip() and line[0].isdigit()): return None

        ohlc = CustomOHLC()
        ohlc.Symbol = config.Symbol

        try:
            data = line.split(',')
            ohlc.Time = datetime.strptime(data[0], "%Y-%m-%d")
            ohlc.EndTime = ohlc.Time + timedelta(days=1)
            ohlc.Value = data[5]
            ohlc["open"] = float(data[1])
            ohlc["high"] = float(data[2])
            ohlc["low"] = float(data[3])
            ohlc["close"] = float(data[4])
        
        except ValueError:
            return None
        
        return ohlc