from .base_analyzer import BaseAnalyzer
import pandas as pd
import numpy as np

class US_Analyzer(BaseAnalyzer):
    def __init__(self):
        super().__init__('US')
        self.mapping = {
            'income': {
                'totalRevenue': 'revenue', 'costOfRevenue': 'cogs', 'grossProfit': 'gross_profit',
                'sellingGeneralAndAdministrative': 'sga_exp', 
                'researchAndDevelopment': 'rd_exp', 'interestExpense': 'fin_exp',
                'incomeBeforeTax': 'total_profit', 'incomeTaxExpense': 'income_tax', 'netIncome': 'net_income',
                'ebit': 'ebit',
            },
            'balance': {
                'cashAndCashEquivalentsAtCarryingValue': 'cash', 'currentNetReceivables': 'receivables', 'inventory': 'inventory',
                'propertyPlantEquipment': 'fixed_assets', 'totalAssets': 'total_assets', 'goodwill': 'goodwill',
                'otherCurrentAssets': 'prepayment', 'longTermInvestments': 'lt_invest', 'otherNonCurrentAssets': 'other_assets',
                'shortTermDebt': 'short_term_debt', 'currentLongTermDebt': 'short_term_debt_part',
                'longTermDebt': 'long_term_debt', 'totalLiabilities': 'total_liabilities', 'totalShareholderEquity': 'total_equity',
                'currentAccountsPayable': 'accounts_payable', 'deferredRevenue': 'adv_receipts',
            },
            'cashflow': {
                'operatingCashflow': 'ocf', 'capitalExpenditures': 'capex', 'dividendPayout': 'dividends'
            }
        }

    def _post_process_columns(self, df, type):
        if type == 'balance':
            if 'short_term_debt' not in df.columns: df['short_term_debt'] = 0
            if 'short_term_debt_part' in df.columns:
                df['short_term_debt'] = df['short_term_debt'].fillna(0) + df['short_term_debt_part'].fillna(0)

        if type == 'income':
            if 'gross_profit' not in df.columns and 'revenue' in df.columns and 'cogs' in df.columns:
                df['gross_profit'] = df['revenue'] - df['cogs']

            if 'ebit' not in df.columns and 'total_profit' in df.columns and 'fin_exp' in df.columns:
                df['ebit'] = df['total_profit'] + df['fin_exp']
        return df

    def calculate_indicators(self, df_merged, market_metrics, historical_metrics):

        def calc_yoy_growth(df, col):
            growths = []
            for idx, row in df.iterrows():
                try:
                    curr_date = str(row['date_str'])
                    target_date = str(int(curr_date) - 10000)
                    
                    match = df[df['date_str'] == target_date]
                    if not match.empty and col in row and pd.notna(row[col]):
                        prev_val = match.iloc[0].get(col)
                        curr_val = row[col]
                        if prev_val and prev_val != 0:
                            growths.append((curr_val - prev_val) / abs(prev_val))
                        else:
                            growths.append(np.nan)
                    else:
                        growths.append(np.nan)
                except:
                    growths.append(np.nan)
            return growths

        if 'revenue' in df_merged.columns:
            df_merged['RevenueGrowth'] = calc_yoy_growth(df_merged, 'revenue')
        if 'net_income' in df_merged.columns:
            df_merged['NetIncomeGrowth'] = calc_yoy_growth(df_merged, 'net_income')
        
        # Margins & Returns
        if 'gross_profit' in df_merged.columns:
            df_merged['GrossMargin'] = self._safe_div(df_merged['gross_profit'], df_merged['revenue'])
        if 'net_income' in df_merged.columns:
            df_merged['NetMargin'] = self._safe_div(df_merged['net_income'], df_merged['revenue'])
            if 'total_equity' in df_merged.columns:
                df_merged['ROE'] = self._safe_div(df_merged['net_income'], df_merged['total_equity'])
            if 'total_assets' in df_merged.columns:
                df_merged['ROA'] = self._safe_div(df_merged['net_income'], df_merged['total_assets'])

        # ROIC
        if 'short_term_debt' in df_merged.columns and 'long_term_debt' in df_merged.columns:
            df_merged['InterestBearingDebt'] = df_merged['short_term_debt'].fillna(0) + df_merged['long_term_debt'].fillna(0)
            if 'total_equity' in df_merged.columns:
                df_merged['InvestedCapital'] = df_merged['total_equity'] + df_merged['InterestBearingDebt']
                if 'ebit' in df_merged.columns:
                    df_merged['ROIC'] = self._safe_div(df_merged['ebit'], df_merged['InvestedCapital'])

        # Cash Flow
        if 'capex' in df_merged.columns:
            df_merged['Capex'] = df_merged['capex'].abs()
            if 'ocf' in df_merged.columns:
                df_merged['FCF'] = df_merged['ocf'] - df_merged['Capex']

        # Expenses
        if 'sga_exp' in df_merged.columns:
             df_merged['SgaRatio'] = self._safe_div(df_merged['sga_exp'], df_merged['revenue'])
        if 'rd_exp' in df_merged.columns:
            df_merged['RDRatio'] = self._safe_div(df_merged['rd_exp'], df_merged['revenue'])
        if 'income_tax' in df_merged.columns and 'total_profit' in df_merged.columns:
            df_merged['TaxRate'] = self._safe_div(df_merged['income_tax'], df_merged['total_profit'])

        # Other Expense Ratio
        if 'GrossMargin' in df_merged.columns and 'NetMargin' in df_merged.columns:
            other_ratio = df_merged['GrossMargin'] - df_merged['NetMargin']
            if 'SgaRatio' in df_merged.columns:
                other_ratio = other_ratio - df_merged['SgaRatio'].fillna(0)
            if 'RDRatio' in df_merged.columns:
                other_ratio = other_ratio - df_merged['RDRatio'].fillna(0)
            df_merged['OtherExpenseRatio'] = other_ratio

        # Depreciation Expense Ratio
        if 'depreciation' in df_merged.columns and 'revenue' in df_merged.columns:
            df_merged['DepreciationRatio'] = self._safe_div(df_merged['depreciation'].abs(), df_merged['revenue'])
        
        if 'total_assets' in df_merged.columns:
            assets = df_merged['total_assets']

            def series_or_zero(column_name: str) -> pd.Series:
                if column_name in df_merged.columns:
                    return pd.to_numeric(df_merged[column_name], errors='coerce').fillna(0)
                return pd.Series(0, index=df_merged.index, dtype='float64')

            # Asset Ratios
            if 'cash' in df_merged.columns: df_merged['CashRatio'] = self._safe_div(df_merged['cash'], assets)
            if 'inventory' in df_merged.columns: df_merged['InventoryRatio'] = self._safe_div(df_merged['inventory'], assets)
            if 'receivables' in df_merged.columns: df_merged['ReceivablesRatio'] = self._safe_div(df_merged['receivables'], assets)
            if 'accounts_payable' in df_merged.columns: df_merged['PayablesRatio'] = self._safe_div(df_merged['accounts_payable'], assets)
            if 'prepayment' in df_merged.columns: df_merged['PrepaymentRatio'] = self._safe_div(df_merged['prepayment'], assets)
            if 'fixed_assets' in df_merged.columns: df_merged['FixedAssetsRatio'] = self._safe_div(df_merged['fixed_assets'], assets)
            if 'lt_invest' in df_merged.columns: df_merged['LongTermInvestmentRatio'] = self._safe_div(df_merged['lt_invest'], assets)
            if 'goodwill' in df_merged.columns: df_merged['GoodwillRatio'] = self._safe_div(df_merged['goodwill'], assets)

            # Other Assets Ratio (as a residual)
            known_assets_ratio = (series_or_zero('CashRatio') + series_or_zero('InventoryRatio') + series_or_zero('ReceivablesRatio') +
                                series_or_zero('PrepaymentRatio') + series_or_zero('FixedAssetsRatio') + series_or_zero('LongTermInvestmentRatio') +
                                series_or_zero('GoodwillRatio'))
            df_merged['OtherAssetsRatio'] = 1 - known_assets_ratio

            # Liability Ratios
            adv = series_or_zero('adv_receipts')
            df_merged['AdvanceReceiptsRatio'] = self._safe_div(adv, assets)

            st_debt = series_or_zero('short_term_debt')
            lt_debt = series_or_zero('long_term_debt')

            df_merged['ShortTermDebtRatio'] = self._safe_div(st_debt, assets)
            df_merged['LongTermDebtRatio'] = self._safe_div(lt_debt, assets)
            df_merged['InterestBearingDebtRatio'] = self._safe_div(st_debt + lt_debt, assets)

            # Operating Assets Ratio
            inv_ratio = series_or_zero('InventoryRatio')
            rec_ratio = series_or_zero('ReceivablesRatio')
            prep_ratio = series_or_zero('PrepaymentRatio')
            pay_ratio = series_or_zero('PayablesRatio')
            adv_ratio = series_or_zero('AdvanceReceiptsRatio')
            df_merged['OperatingAssetsRatio'] = inv_ratio + rec_ratio + prep_ratio - pay_ratio - adv_ratio

        # Turnover Days
        if 'cogs' in df_merged.columns:
            if 'inventory' in df_merged.columns: 
                df_merged['InventoryDays'] = self._safe_div(df_merged['inventory'] * 365, df_merged['cogs'])
            if 'accounts_payable' in df_merged.columns:
                df_merged['PayablesDays'] = self._safe_div(df_merged['accounts_payable'] * 365, df_merged['cogs'])
        if 'revenue' in df_merged.columns:
            if 'receivables' in df_merged.columns:
                df_merged['ReceivablesDays'] = self._safe_div(df_merged['receivables'] * 365, df_merged['revenue'])
            if 'fixed_assets' in df_merged.columns:
                df_merged['FixedAssetsTurnover'] = self._safe_div(df_merged['revenue'], df_merged['fixed_assets'])
            if 'total_assets' in df_merged.columns:
                df_merged['TotalAssetTurnover'] = self._safe_div(df_merged['revenue'], df_merged['total_assets'])

        # Market Metrics
        if not df_merged.empty:
            df_merged.loc[df_merged.index[0], 'Price'] = market_metrics.get('price')
            df_merged.loc[df_merged.index[0], 'MarketCap'] = market_metrics.get('market_cap')
            df_merged.loc[df_merged.index[0], 'PE'] = market_metrics.get('pe')
            df_merged.loc[df_merged.index[0], 'PB'] = market_metrics.get('pb')

        # US Market doesn't typically have per-employee metrics readily available in the same way

        return df_merged