"""
ml.forecast — ML-based financial forecasting using XGBoost + feature engineering.

Models trained on historical financials + macro factors to produce
revenue, margin, and FCF projections with confidence intervals.
"""
from __future__ import annotations

import warnings
import numpy as np
import pandas as pd
from dataclasses import dataclass, field

warnings.filterwarnings("ignore")


@dataclass
class ForecastResult:
    """Output of a forecast run."""
    metric: str
    years: list[int]
    point: list[float]             # point estimates
    lower: list[float]             # 80% CI lower
    upper: list[float]             # 80% CI upper
    cagr: float                    # implied CAGR
    drivers: dict[str, float]      # feature importances
    model_name: str = "XGBoost"

    def to_df(self) -> pd.DataFrame:
        return pd.DataFrame({
            "year": self.years,
            "forecast": self.point,
            "lower_80": self.lower,
            "upper_80": self.upper,
        }).set_index("year")

    def summary(self):
        from finverse.utils.display import console, fmt_currency, fmt_pct
        from rich.table import Table
        from rich import box

        table = Table(title=f"{self.metric} Forecast ({self.model_name})", box=box.SIMPLE_HEAD, header_style="bold blue")
        table.add_column("Year")
        table.add_column("Forecast", justify="right")
        table.add_column("80% CI", justify="right")

        for yr, pt, lo, hi in zip(self.years, self.point, self.lower, self.upper):
            table.add_row(
                str(yr),
                fmt_currency(pt),
                f"{fmt_currency(lo)} – {fmt_currency(hi)}",
            )

        console.print(table)
        console.print(f"  Implied CAGR: [bold]{fmt_pct(self.cagr)}[/bold]")

        if self.drivers:
            console.print("\n  [dim]Key drivers:[/dim]")
            sorted_drivers = sorted(self.drivers.items(), key=lambda x: -x[1])
            for name, imp in sorted_drivers[:5]:
                bar = "█" * int(imp * 20)
                console.print(f"    {name:<25} {bar} {fmt_pct(imp)}")
        console.print()


def _build_features(series: pd.Series, macro_df: pd.DataFrame | None = None) -> pd.DataFrame:
    """Engineer features from a financial time series."""
    df = pd.DataFrame({"value": series.values}, index=range(len(series)))

    df["lag1"] = df["value"].shift(1)
    df["lag2"] = df["value"].shift(2)
    df["lag3"] = df["value"].shift(3)
    df["yoy_growth"] = df["value"].pct_change()
    df["yoy_growth_lag1"] = df["yoy_growth"].shift(1)
    df["rolling_mean_3"] = df["value"].rolling(3).mean()
    df["rolling_std_3"] = df["value"].rolling(3).std()
    df["trend"] = range(len(df))
    df["trend_sq"] = df["trend"] ** 2

    if macro_df is not None and not macro_df.empty:
        for col in macro_df.columns[:3]:
            try:
                aligned = macro_df[col].resample("YE").last().values
                if len(aligned) >= len(df):
                    df[f"macro_{col}"] = aligned[:len(df)]
            except Exception:
                pass

    return df.dropna()


def _xgboost_forecast(
    series: pd.Series,
    n_forward: int,
    macro_df: pd.DataFrame | None,
    n_bootstrap: int = 50,
) -> tuple[list[float], list[float], list[float], dict]:
    """Core XGBoost forecasting with bootstrap confidence intervals."""
    try:
        from xgboost import XGBRegressor
    except ImportError:
        raise ImportError("xgboost required: pip install xgboost")

    if len(series) < 4:
        last = series.iloc[-1] if len(series) > 0 else 0
        hist_cagr = series.pct_change().mean() if len(series) > 1 else 0.05
        pts = [last * (1 + hist_cagr) ** i for i in range(1, n_forward + 1)]
        std = abs(last) * 0.1
        return pts, [p - std for p in pts], [p + std for p in pts], {}

    features_df = _build_features(series, macro_df)
    if features_df.empty or len(features_df) < 3:
        last = series.iloc[-1]
        cagr = series.pct_change().mean()
        pts = [last * (1 + cagr) ** i for i in range(1, n_forward + 1)]
        return pts, pts, pts, {}

    feature_cols = [c for c in features_df.columns if c != "value"]
    X = features_df[feature_cols].values
    y = features_df["value"].values

    model = XGBRegressor(
        n_estimators=100,
        max_depth=3,
        learning_rate=0.1,
        subsample=0.8,
        random_state=42,
        verbosity=0,
    )
    model.fit(X, y)

    importances = dict(zip(feature_cols, model.feature_importances_))

    all_preds = []
    for _ in range(n_bootstrap):
        idx = np.random.choice(len(X), size=len(X), replace=True)
        m = XGBRegressor(
            n_estimators=80, max_depth=3, learning_rate=0.1,
            subsample=0.8, random_state=np.random.randint(1000), verbosity=0
        )
        m.fit(X[idx], y[idx])

        current_series = list(series.values)
        preds = []
        for step in range(n_forward):
            feat_series = pd.Series(current_series)
            feat_df = _build_features(feat_series, macro_df)
            if feat_df.empty:
                preds.append(current_series[-1] * (1 + 0.05))
            else:
                last_feat = feat_df[feature_cols].iloc[-1:].values
                last_feat = np.nan_to_num(last_feat)
                pred = float(m.predict(last_feat)[0])
                preds.append(pred)
                current_series.append(pred)
        all_preds.append(preds)

    all_preds = np.array(all_preds)
    point = list(np.median(all_preds, axis=0))
    lower = list(np.percentile(all_preds, 10, axis=0))
    upper = list(np.percentile(all_preds, 90, axis=0))

    return point, lower, upper, importances


def revenue(
    data,
    n_years: int = 5,
    macro_df: pd.DataFrame | None = None,
) -> ForecastResult:
    """
    Forecast revenue using ML on historical financials + macro factors.

    Parameters
    ----------
    data     : TickerData or pd.Series of historical revenue ($B)
    n_years  : int — number of years to forecast (default 5)
    macro_df : optional FRED macro DataFrame to use as features

    Returns
    -------
    ForecastResult

    Example
    -------
    >>> from finverse import pull
    >>> from finverse.ml import forecast
    >>> data = pull.ticker("AAPL")
    >>> fc = forecast.revenue(data, n_years=5)
    >>> fc.summary()
    """
    from finverse.utils.display import console

    if hasattr(data, "revenue_history"):
        series = data.revenue_history
        name = getattr(data, "name", "Company")
    elif isinstance(data, pd.Series):
        series = data
        name = data.name or "Revenue"
    else:
        raise TypeError("data must be a TickerData or pd.Series")

    series = series.dropna().sort_index()
    console.print(f"[dim]Forecasting revenue ({len(series)} years of history)...[/dim]")

    last_year = series.index[-1]
    last_year_int = last_year.year if hasattr(last_year, "year") else int(str(last_year)[:4])
    forecast_years = list(range(last_year_int + 1, last_year_int + n_years + 1))

    point, lower, upper, importances = _xgboost_forecast(series, n_years, macro_df)

    last_val = series.iloc[-1]
    final_val = point[-1] if point else last_val
    cagr = (final_val / last_val) ** (1 / n_years) - 1 if last_val != 0 else 0

    console.print(f"[green]✓[/green] Revenue forecast complete — implied CAGR {cagr*100:.1f}%")

    return ForecastResult(
        metric="Revenue ($B)",
        years=forecast_years,
        point=point,
        lower=lower,
        upper=upper,
        cagr=cagr,
        drivers=importances,
    )


def margins(
    data,
    margin_type: str = "ebitda",
    n_years: int = 5,
    macro_df: pd.DataFrame | None = None,
) -> ForecastResult:
    """
    Forecast EBITDA or net margins using ML.

    Parameters
    ----------
    data        : TickerData
    margin_type : "ebitda" or "net" (default "ebitda")
    n_years     : int
    macro_df    : optional macro DataFrame

    Returns
    -------
    ForecastResult (values are margin %, e.g. 0.32 = 32%)

    Example
    -------
    >>> fc = forecast.margins(data, margin_type="ebitda")
    >>> fc.summary()
    """
    from finverse.utils.display import console

    if not hasattr(data, "revenue_history"):
        raise TypeError("data must be a TickerData object")

    rev = data.revenue_history
    if margin_type == "ebitda":
        num = data.ebitda_history
        label = "EBITDA Margin"
    else:
        num = data.net_income_history
        label = "Net Margin"

    common_idx = rev.index.intersection(num.index)
    if len(common_idx) < 3:
        console.print(f"[yellow]Warning: insufficient data for margin forecast[/yellow]")
        dummy = pd.Series([0.25] * 3)
    else:
        margin_series = (num.loc[common_idx] / rev.loc[common_idx]).dropna()
        margin_series = margin_series[margin_series.between(-1, 2)]

    console.print(f"[dim]Forecasting {label}...[/dim]")

    last_year_int = common_idx[-1].year if hasattr(common_idx[-1], "year") else int(str(common_idx[-1])[:4])
    forecast_years = list(range(last_year_int + 1, last_year_int + n_years + 1))

    point, lower, upper, importances = _xgboost_forecast(margin_series, n_years, macro_df)

    last_val = margin_series.iloc[-1]
    final_val = point[-1] if point else last_val
    cagr = (final_val - last_val) / n_years

    console.print(f"[green]✓[/green] {label} forecast complete")

    return ForecastResult(
        metric=label,
        years=forecast_years,
        point=point,
        lower=lower,
        upper=upper,
        cagr=cagr,
        drivers=importances,
    )


def wacc(
    data,
    risk_free_rate: float | None = None,
    market_premium: float = 0.055,
) -> dict:
    """
    Estimate WACC from market data — no manual inputs required.

    Uses beta from price history, credit spread proxy from debt/equity,
    and a risk-free rate from current 10Y treasury (via FRED if available).

    Parameters
    ----------
    data            : TickerData
    risk_free_rate  : float — override (e.g. 0.045). If None, uses 10Y treasury.
    market_premium  : float — equity risk premium (default 5.5%)

    Returns
    -------
    dict with keys: wacc, cost_of_equity, cost_of_debt, tax_rate, weights

    Example
    -------
    >>> est = forecast.wacc(data)
    >>> print(f"WACC: {est['wacc']:.1%}")
    """
    from finverse.utils.display import console

    console.print("[dim]Estimating WACC from market data...[/dim]")

    result = {
        "wacc": 0.095,
        "cost_of_equity": 0.10,
        "cost_of_debt": 0.045,
        "tax_rate": 0.21,
        "weights": {"equity": 0.8, "debt": 0.2},
        "beta": 1.0,
        "risk_free_rate": risk_free_rate or 0.045,
        "method": "estimated",
    }

    try:
        if not data.price_history.empty and not data.price_history["Close"].empty:
            try:
                import yfinance as yf
                spy = yf.Ticker("SPY").history(period="3y")["Close"].pct_change().dropna()
                stock_ret = data.price_history["Close"].pct_change().dropna()
                common = spy.index.intersection(stock_ret.index)
                if len(common) > 60:
                    beta = np.cov(stock_ret.loc[common], spy.loc[common])[0][1] / np.var(spy.loc[common])
                    beta = float(np.clip(beta, 0.3, 3.0))
                    result["beta"] = beta
            except Exception:
                pass

        rf = risk_free_rate or 0.045
        result["risk_free_rate"] = rf

        beta = result["beta"]
        ke = rf + beta * market_premium
        result["cost_of_equity"] = round(ke, 4)

        debt = data.total_debt or 0
        cash = data.cash or 0
        net_debt = max(debt - cash, 0)
        mkt_cap = (data.market_cap or 1e12) / 1e9

        total_capital = mkt_cap + net_debt
        w_e = mkt_cap / total_capital if total_capital > 0 else 0.8
        w_d = net_debt / total_capital if total_capital > 0 else 0.2

        result["weights"] = {"equity": round(w_e, 3), "debt": round(w_d, 3)}

        kd = result["cost_of_debt"]
        tax = result["tax_rate"]
        wacc_val = w_e * ke + w_d * kd * (1 - tax)
        result["wacc"] = round(wacc_val, 4)
        result["method"] = "market-derived"

        console.print(
            f"[green]✓[/green] WACC = {wacc_val:.1%} "
            f"(β={beta:.2f}, Ke={ke:.1%}, Kd={kd:.1%}, "
            f"E:{w_e:.0%} / D:{w_d:.0%})"
        )

    except Exception as e:
        console.print(f"[yellow]Warning: WACC estimation error ({e}), using defaults[/yellow]")

    return result
