# region imports
from AlgorithmImports import *

from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans

from industry_codes import industry_codes
# endregion

fundamental_factors = [
    'operation_ratios.revenue_growth.three_months',              # How a company's sales are affected over time, possibly reflecting policy impacts like tax changes or sector-specific subsidies.
    #'valuation_ratios.first_year_estimated_eps_growth',         # Indicates how profitable a company has become, potentially reflecting beneficial or adverse policy effects.
    'valuation_ratios.pe_ratio',                                 # Changes in this ratio might reflect investor expectations under different political climates. For instance, sectors expected to benefit from new policies might see an increase in their P/E ratios due to higher anticipated earnings.
    'operation_ratios.roe.three_months',                         # Indicate how well companies utilize their assets to generate profit, which could be influenced by corporate tax policies or regulatory environments.
    'operation_ratios.roa.three_months',                         # Indicate how well companies utilize their assets to generate profit, which could be influenced by corporate tax policies or regulatory environments.
    'operation_ratios.current_ratio.three_months',               # To see how policies affecting economic stability might impact a company's short-term financial health.
    'operation_ratios.total_debt_equity_ratio.three_months',     # Policies affecting interest rates or corporate taxes could influence a company's capital structure.
    'valuation_ratios.trailing_dividend_yield',                  # Reflects income generation for investors, which might be influenced by tax policies on dividends.
]

class CreativeRedCow(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2012, 3, 1)
        self.set_end_date(2024, 9, 30)
        self.set_cash(10_000_000)
        self.set_warm_up(self.start_date - datetime(1998, 2, 1))

        self.universe_settings.schedule.on(self.date_rules.month_start())
        self.add_universe(self._select_assets)

        self._daily_returns_by_industry = {industry: pd.Series() for industry in industry_codes}
        
        self._scaler = StandardScaler()
        self._clustering_model = KMeans(n_clusters=3, random_state=0)
        self._lookback = timedelta(12*365) # 144 months = 12 years = 3 presidential terms.
        self._previous_selection_time = None
        self._winning_party_by_result_date = {date(1996, 11, 6): "D", date(2000, 11, 8): "R", date(2004, 11, 3): "R", date(2008, 11, 5): "D", date(2012, 11, 7): "D", date(2016, 11, 9): "R", date(2020, 11, 8): "D"}
        self._periods_by_party = {
            'D': [
                (datetime(1996, 11, 6), datetime(2000, 11, 7)), 
                (datetime(2008, 11, 5), datetime(2016, 11, 8)),
                (datetime(2020, 11, 8), self.end_date)
            ],
            'R': [
                (datetime(2000, 11, 8), datetime(2008, 11, 4)),
                (datetime(2016, 11, 9), datetime(2020, 11, 7))
            ]
        }

    def _has_all_factors(self, fundamental):
        values = []
        for factor in fundamental_factors:
            values.append(np.isnan(eval(f"fundamental.{factor}")))
            #values.append(eval(f"fundamental.{factor}"))
        return not any(values)

    def _select_assets(self, fundamentals):
        ruling_party = [party for t, party in self._winning_party_by_result_date.items() if t < self.time.date()][-1]

        # Update the daily returns of each industry.
        if self._previous_selection_time:
            # Get the daily returns of each industry leader.
            daily_returns = self.history([leader['symbol'] for leader in self._leader_by_industry.values()], self._previous_selection_time-timedelta(2), self.time, Resolution.DAILY)['close'].unstack(0).pct_change()
            for industry in industry_codes:
                if industry not in self._leader_by_industry or self._leader_by_industry[industry]['symbol'] not in daily_returns.columns:
                    continue # Fill with 0 later
                # Append the daily returns of the current leader to the series of daily returns for the industry.
                industry_daily_returns = pd.concat([self._daily_returns_by_industry[industry], daily_returns[self._leader_by_industry[industry]['symbol']].dropna()])
                self._daily_returns_by_industry[industry] = industry_daily_returns[~industry_daily_returns.index.duplicated(keep='first')]
                # Trim off history that has fallen out of the lookback window.
                self._daily_returns_by_industry[industry] = self._daily_returns_by_industry[industry][self._daily_returns_by_industry[industry].index >= self.time - self._lookback]
        self._previous_selection_time = self.time

        # Get the largest asset of each industry (and its fundamental factors).
        self._leader_by_industry = {}
        for industry_code in industry_codes:
            industry_assets = [
                f for f in fundamentals 
                if f.asset_classification.morningstar_industry_code == eval(f"MorningstarIndustryCode.{industry_code}") and f.market_cap and self._has_all_factors(f)
            ]
            if industry_assets:
                leader = sorted(industry_assets, key=lambda f: f.market_cap)[-1]
                leader_factors = {}
                for factor in fundamental_factors + ['symbol']:
                    leader_factors[factor] = eval(f"leader.{factor}") 
                self._leader_by_industry[industry_code] = leader_factors

        # During warm-up, keep the universe empty.
        if self.is_warming_up or not (self.time.year % 4 == 0 and self.time.month == 11):
            return []

        # Cluster the assets based on their fundamental factors.
        factors_df = pd.DataFrame(columns=fundamental_factors)
        for industry, leader in self._leader_by_industry.items():
            factors_df.loc[industry, :] = [leader[f] for f in fundamental_factors]
        cluster_by_industry = pd.Series(
            self._clustering_model.fit_predict(self._scaler.fit_transform(factors_df)),
            index=self._leader_by_industry.keys()
        )
        
        # Calculate the performance of each cluster.
        other_party = 'D' if ruling_party == 'R' else 'R'
        cluster_scores = []
        for cluster in cluster_by_industry.unique():
            result = {}

            # Get the daily returns of the cluster.
            industries = cluster_by_industry[cluster_by_industry == cluster].index
            cluster_daily_returns = pd.concat([self._daily_returns_by_industry[industry] for industry in industries], axis=1)
            
            # Select the periods of time when the current party was in office.
            sliced_daily_returns = pd.DataFrame()
            for start_date, end_date in self._periods_by_party[ruling_party]:
                sliced_daily_returns = pd.concat([sliced_daily_returns, cluster_daily_returns.loc[start_date:end_date]])
            
            # Calculate metric (ex: Return of an equal-weighted portfolio).
            if sliced_daily_returns.empty:
                result[ruling_party] = 0
            else:
                equity_curve = (sliced_daily_returns.fillna(0).mean(axis=1) + 1).cumprod()
                result[ruling_party] = (equity_curve.iloc[-1] - equity_curve.iloc[0]) / equity_curve.iloc[0]

            # Select the periods of time when the other party was in office.
            sliced_daily_returns = pd.DataFrame()
            for start_date, end_date in self._periods_by_party[other_party]:
                sliced_daily_returns = pd.concat([sliced_daily_returns, cluster_daily_returns.loc[start_date:end_date]])
            
            # Calculate metric (ex: Return of an equal-weighted portfolio).
            if sliced_daily_returns.empty:
                result[other_party] = 0
            else:
                equity_curve = (sliced_daily_returns.fillna(0).mean(axis=1) + 1).cumprod()
                result[other_party] = (equity_curve.iloc[-1] - equity_curve.iloc[0]) / equity_curve.iloc[0]

            # Select the ROC over the trailing month.
            result['last_month'] = (cluster_daily_returns.fillna(0).mean(axis=1) + 1).cumprod().pct_change(21).iloc[-1]
            
            cluster_scores.append(result)

        # Whichever cluster performs best over the last month, find which party it leans towards.
        winning_cluster = sorted(cluster_scores, key=lambda scores: scores['last_month'])[-1]
        predicted_party = ruling_party if winning_cluster[ruling_party] > winning_cluster[other_party] else other_party
        self.debug(f'Predicted party: {predicted_party}')
        self.log(f'Predicted party: {predicted_party}')

        # Return an empty universe
        return []
        

# region imports
from AlgorithmImports import *
# endregion

industry_codes = [
    'AGRICULTURAL_INPUTS',
    'BUILDING_MATERIALS',
    'CHEMICALS',
    'SPECIALTY_CHEMICALS',
    'LUMBER_AND_WOOD_PRODUCTION',
    'PAPER_AND_PAPER_PRODUCTS',
    'ALUMINUM',
    'COPPER',
    'OTHER_INDUSTRIAL_METALS_AND_MINING',
    'GOLD',
    'SILVER',
    'OTHER_PRECIOUS_METALS_AND_MINING',
    'COKING_COAL',
    'STEEL',
    'AUTO_AND_TRUCK_DEALERSHIPS',
    'AUTO_MANUFACTURERS',
    'AUTO_PARTS',
    'RECREATIONAL_VEHICLES',
    'FURNISHINGS',
    'FIXTURES_AND_APPLIANCES',
    'RESIDENTIAL_CONSTRUCTION',
    'TEXTILE_MANUFACTURING',
    'APPAREL_MANUFACTURING',
    'FOOTWEAR_AND_ACCESSORIES',
    'PACKAGING_AND_CONTAINERS',
    'PERSONAL_SERVICES',
    'RESTAURANTS',
    'APPAREL_RETAIL',
    'DEPARTMENT_STORES',
    'HOME_IMPROVEMENT_RETAIL',
    'LUXURY_GOODS',
    'INTERNET_RETAIL',
    'SPECIALTY_RETAIL',
    'GAMBLING',
    'LEISURE',
    'LODGING',
    'RESORTS_AND_CASINOS',
    'TRAVEL_SERVICES',
    'ASSET_MANAGEMENT',
    'BANKS_DIVERSIFIED',
    'BANKS_REGIONAL',
    'MORTGAGE_FINANCE',
    'CAPITAL_MARKETS',
    'FINANCIAL_DATA_AND_STOCK_EXCHANGES',
    'INSURANCE_LIFE',
    'INSURANCE_PROPERTY_AND_CASUALTY',
    'INSURANCE_REINSURANCE',
    'INSURANCE_SPECIALTY',
    'INSURANCE_BROKERS',
    'INSURANCE_DIVERSIFIED',
    'SHELL_COMPANIES',
    'FINANCIAL_CONGLOMERATES',
    'CREDIT_SERVICES',
    'REAL_ESTATE_DEVELOPMENT',
    'REAL_ESTATE_SERVICES',
    'REAL_ESTATE_DIVERSIFIED',
    'REIT_HEALTHCARE_FACILITIES',
    'REIT_HOTEL_AND_MOTEL',
    'REIT_INDUSTRIAL',
    'REIT_OFFICE',
    'REIT_RESIDENTIAL',
    'REIT_RETAIL',
    'REIT_MORTGAGE',
    'REIT_SPECIALTY',
    'REIT_DIVERSIFIED',
    'BEVERAGES_BREWERS',
    'BEVERAGES_WINERIES_AND_DISTILLERIES',
    'BEVERAGES_NON_ALCOHOLIC',
    'CONFECTIONERS',
    'FARM_PRODUCTS',
    'HOUSEHOLD_AND_PERSONAL_PRODUCTS',
    'PACKAGED_FOODS',
    'EDUCATION_AND_TRAINING_SERVICES',
    'DISCOUNT_STORES',
    'FOOD_DISTRIBUTION',
    'GROCERY_STORES',
    'TOBACCO',
    'BIOTECHNOLOGY',
    'DRUG_MANUFACTURERS_GENERAL',
    'DRUG_MANUFACTURERS_SPECIALTY_AND_GENERIC',
    'HEALTHCARE_PLANS',
    'MEDICAL_CARE_FACILITIES',
    'PHARMACEUTICAL_RETAILERS',
    'HEALTH_INFORMATION_SERVICES',
    'MEDICAL_DEVICES',
    'MEDICAL_INSTRUMENTS_AND_SUPPLIES',
    'DIAGNOSTICS_AND_RESEARCH',
    'MEDICAL_DISTRIBUTION',
    'UTILITIES_INDEPENDENT_POWER_PRODUCERS',
    'UTILITIES_RENEWABLE',
    'UTILITIES_REGULATED_WATER',
    'UTILITIES_REGULATED_ELECTRIC',
    'UTILITIES_REGULATED_GAS',
    'UTILITIES_DIVERSIFIED',
    'TELECOM_SERVICES',
    'ADVERTISING_AGENCIES',
    'PUBLISHING',
    'BROADCASTING',
    'ENTERTAINMENT',
    'INTERNET_CONTENT_AND_INFORMATION',
    'ELECTRONIC_GAMING_AND_MULTIMEDIA',
    'OIL_AND_GAS_DRILLING',
    'OIL_AND_GAS_E_AND_P',
    'OIL_AND_GAS_INTEGRATED',
    'OIL_AND_GAS_MIDSTREAM',
    'OIL_AND_GAS_REFINING_AND_MARKETING',
    'OIL_AND_GAS_EQUIPMENT_AND_SERVICES',
    'THERMAL_COAL',
    'URANIUM',
    'AEROSPACE_AND_DEFENSE',
    'SPECIALTY_BUSINESS_SERVICES',
    'CONSULTING_SERVICES',
    'RENTAL_AND_LEASING_SERVICES',
    'SECURITY_AND_PROTECTION_SERVICES',
    'STAFFING_AND_EMPLOYMENT_SERVICES',
    'CONGLOMERATES',
    'ENGINEERING_AND_CONSTRUCTION',
    'INFRASTRUCTURE_OPERATIONS',
    'BUILDING_PRODUCTS_AND_EQUIPMENT',
    'FARM_AND_HEAVY_CONSTRUCTION_MACHINERY',
    'INDUSTRIAL_DISTRIBUTION',
    'BUSINESS_EQUIPMENT_AND_SUPPLIES',
    'SPECIALTY_INDUSTRIAL_MACHINERY',
    'METAL_FABRICATION',
    'POLLUTION_AND_TREATMENT_CONTROLS',
    'TOOLS_AND_ACCESSORIES',
    'ELECTRICAL_EQUIPMENT_AND_PARTS',
    'AIRPORTS_AND_AIR_SERVICES',
    'AIRLINES',
    'RAILROADS',
    'MARINE_SHIPPING',
    'TRUCKING',
    'INTEGRATED_FREIGHT_AND_LOGISTICS',
    'WASTE_MANAGEMENT',
    'INFORMATION_TECHNOLOGY_SERVICES',
    'SOFTWARE_APPLICATION',
    'SOFTWARE_INFRASTRUCTURE',
    'COMMUNICATION_EQUIPMENT',
    'COMPUTER_HARDWARE',
    'CONSUMER_ELECTRONICS',
    'ELECTRONIC_COMPONENTS',
    'ELECTRONICS_AND_COMPUTER_DISTRIBUTION',
    'SCIENTIFIC_AND_TECHNICAL_INSTRUMENTS',
    'SEMICONDUCTOR_EQUIPMENT_AND_MATERIALS',
    'SEMICONDUCTORS',
    'SOLAR'
]