# region imports
from AlgorithmImports import *
# endregion
from datetime import datetime, timedelta
import xml.etree.ElementTree as ET
from System.Net import WebRequest
from QuantConnect.Data import BaseData

class FredM2Data(BaseData):
    """Custom data class for importing M2 Money Supply data from FRED"""
    
    def __init__(self):
        self.Symbol = None
        self.Value = 0
        self.Time = datetime.min
        
    def GetSource(self, config, date, isLiveMode):
        """Return the URL for downloading the data"""
        # Replace YOUR_API_KEY with your actual FRED API key
        api_key = "f156cff53e458a5f18a8f421ebb944c1"
        
        # Build the FRED API endpoint URL
        # M2SL is the series ID for M2 Money Stock, Seasonally Adjusted
        url = f"https://api.stlouisfed.org/fred/series/observations?series_id=M2SL&api_key={api_key}&file_type=xml"
        
        # For backtesting, we want to get all historical data at once to avoid multiple API calls
        return url
        
    def Reader(self, config, line, date, isLiveMode):
        """Parse the XML data from FRED"""
        if line == "" or line is None:
            return None
            
        # Create a new instance
        data = FredM2Data()
        data.Symbol = config.Symbol
        
        try:
            # Parse the XML string
            root = ET.fromstring(line)
            
            # Find the most recent observation before our requested date
            observations = root.findall(".//observation")
            
            closest_observation = None
            for observation in observations:
                obs_date = datetime.strptime(observation.get('date'), "%Y-%m-%d")
                
                # Skip future dates
                if obs_date > date:
                    continue
                
                # If this is the first valid observation or it's more recent than our current closest
                if closest_observation is None or obs_date > closest_observation[0]:
                    closest_observation = (obs_date, observation)
            
            if closest_observation is not None:
                value_str = closest_observation[1].get('value')
                
                # Check if the value is available (not "." which FRED uses for missing values)
                if value_str != ".":
                    # Convert value to float (it's in billions of dollars)
                    data.Value = float(value_str)
                    data.Time = closest_observation[0]
                    return data
            
            return None
        
        except Exception as e:
            # Log the error and return None for unsuccessful parsing
            print(f"Error parsing FRED M2 data: {e}")
            return None

# region imports
from AlgorithmImports import *
from datetime import datetime, timedelta
import json
from System.Net import WebRequest

#from FredM2Data import FredM2Data
# endregion
class M2AdjustedNasdaqAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 12, 31)
        self.SetCash(100000)
        
        # Add Nasdaq index data (using QQQ as a proxy)
        self.nasdaq = self.AddEquity("QQQ", Resolution.Daily).Symbol
        
        # Add custom M2 data from FRED
        self.m2 = self.add_data(FredM2Data, "M2SL").Symbol
        
        
        # Set benchmark to Nasdaq
        self.SetBenchmark(self.nasdaq)
        
        # Dictionary to store M2 data for each month
        self.m2_data = {}
        
        # Store initial values for normalization
        self.initial_nasdaq = None
        self.initial_m2 = None
        
        # Dictionary to store calculated values
        self.m2_adjusted_values = {}
        
        # Flag to track if we've done initial calculations
        self.initialized = False
        
        # Schedule monthly calculations
        self.Schedule.On(self.DateRules.MonthStart(), self.TimeRules.At(9, 30), self.CalculateM2AdjustedIndex)
        
        self.m2_counter = 0  # Add a counter
    
    def OnData(self, data):
        """OnData event is the primary entry point for your algorithm."""
        if self.m2 in data:
            self.m2_counter += 1
            m2_value = data[self.m2].Value
            if m2_value > 0:  # Make sure we have valid data
                month_key = f"{self.Time.year}-{self.Time.month}"
                self.m2_data[month_key] = m2_value
                self.Debug(f"M2 data received for {month_key}: {m2_value} billion (Count: {self.m2_counter})")
        
        # Initialize values once we have both Nasdaq and M2 data
        if not self.initialized and self.nasdaq in data and len(self.m2_data) > 0:
            self.initial_nasdaq = data[self.nasdaq].Price
            # Use the last available M2 data point as the baseline
            self.initial_m2 = list(self.m2_data.values())[-1]
            self.initialized = True
            self.Log(f"Initial values set - Nasdaq: {self.initial_nasdaq}, M2: {self.initial_m2} billion")
        
        # Only trade if we're initialized
        if self.initialized and not self.Portfolio.Invested:
            self.SetHoldings(self.nasdaq, 1.0)
    
    def CalculateM2AdjustedIndex(self):
        """Calculate the M2-adjusted Nasdaq index value on a monthly basis"""
        if not self.initialized:
            return
        
        # Get latest data
        nasdaq_price = self.Securities[self.nasdaq].Price
        
        # Get current month M2 data
        month_key = f"{self.Time.year}-{self.Time.month}"
        if month_key not in self.m2_data:
            # Use the most recent available M2 data
            if len(self.m2_data) > 0:
                m2_value = list(self.m2_data.values())[-1]
            else:
                return  # No M2 data available yet
        else:
            m2_value = self.m2_data[month_key]
        
        # Calculate M2-adjusted Nasdaq index
        # Formula: (Current Nasdaq / Initial Nasdaq) / (Current M2 / Initial M2) * Initial Nasdaq
        # This effectively deflates the Nasdaq by the M2 money supply growth
        m2_ratio = self.initial_m2 / m2_value
        adjusted_nasdaq = nasdaq_price * m2_ratio
        
        # Store the calculated value
        self.m2_adjusted_values[self.Time] = adjusted_nasdaq
        
        # Update chart
        #self.Plot("Nasdaq", "Nasdaq", nasdaq_price)
        self.Plot("M2 Money Supply (Billions)", "M2 Money Supply (Billions)", m2_value)
        self.Plot("M2-Adjusted Nasdaq", "M2-Adjusted Nasdaq", adjusted_nasdaq)
        
        # Log values
        self.Log(f"Date: {self.Time}, Nasdaq: {nasdaq_price}, M2: {m2_value}, M2-Adjusted Nasdaq: {adjusted_nasdaq}")
        
        # Show the inflation-adjusted performance comparison
        real_performance = (adjusted_nasdaq / self.initial_nasdaq - 1) * 100
        nominal_performance = (nasdaq_price / self.initial_nasdaq - 1) * 100
        monetary_effect = nominal_performance - real_performance
        
        self.Log(f"Nasdaq nominal growth: {nominal_performance:.2f}%, M2-adjusted growth: {real_performance:.2f}%")
        self.Log(f"Monetary expansion effect: {monetary_effect:.2f}%")


class FredM2Data(PythonData):
    """Custom data class for importing M2 Money Supply data from FRED using JSON"""

    def __init__(self):
        self.Value = 0
        self.Time = datetime.min
        # Dictionary to cache parsed observation data
        self.observations = {}

            
    def get_source(self,
         config: SubscriptionDataConfig,
         date: datetime,
         is_live: bool) -> SubscriptionDataSource:

        """Return the URL for downloading the JSON data from FRED"""
        # Replace YOUR_API_KEY with your actual FRED API key
        api_key = "f156cff53e458a5f18a8f421ebb944c1"
        # Append file_type=json to receive JSON data
        url = f"https://api.stlouisfed.org/fred/series/observations?series_id=M2SL&api_key={api_key}&file_type=json"
        return SubscriptionDataSource(url, SubscriptionTransportMedium.REMOTE_FILE)
        


    
    
    def Reader(self, config, line, date, isLiveMode):
        """Parse the JSON data from FRED and return data for the specific date"""
        if line == "" or line is None:
            return None
        
        try:
            # If we haven't parsed the JSON data yet
            if not self.observations:
                # Parse the JSON string
                data_dict = json.loads(line)
                
                # Extract all observations and store in a dictionary by date
                for observation in data_dict.get('observations', []):
                    obs_date_str = observation.get('date')
                    value_str = observation.get('value')
                    
                    # Skip missing values
                    if value_str == ".":
                        continue
                    
                    obs_date = datetime.strptime(obs_date_str, "%Y-%m-%d")
                    # Convert to QuantConnect DateTime format
                    qc_date = DateTime(obs_date.year, obs_date.month, obs_date.day)
                    # Store the observation
                    self.observations[qc_date] = float(value_str)
                
                # Log how many observations were parsed
                print(f"Parsed {len(self.observations)} FRED M2 observations from JSON")
            
            # Create data points for all observations that match our date
            # Convert Python datetime to QuantConnect DateTime
            qc_date = DateTime(date.year, date.month, date.day)
            
            # For daily resolution, we need to forward-fill data
            # Find the most recent observation on or before the requested date
            valid_dates = [d for d in self.observations.keys() if d <= qc_date]
            
            if not valid_dates:
                return None
                
            # Get the most recent date
            most_recent_date = max(valid_dates)
            value = self.observations[most_recent_date]
            
            # Create data point
            data = FredM2Data()
            data.Symbol = config.Symbol
            data.Time = date
            data.Value = value
            
            return data
            
        except Exception as e:
            # Log the error and return None for unsuccessful parsing
            print(f"Error parsing FRED M2 JSON data: {e}")
            return None
    
    def Clone(self, fillForward):
        """Clone method required by QuantConnect"""
        return self.MemberwiseClone()