import matplotlib
import numpy as np
import pandas as pd

matplotlib.use("Agg")
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from sklearn.linear_model import LinearRegression

# ── Try SB3 import ─────────────────────────────────────────────
try:
    from stable_baselines3 import PPO

    SB3_AVAILABLE = True
except ImportError:
    SB3_AVAILABLE = False

try:
    from trading_env import TradingEnv

    ENV_AVAILABLE = True
except ImportError:
    ENV_AVAILABLE = False

from sklearn.preprocessing import MinMaxScaler

#  PRICE TREND FORECAST
def forecast_price_trend(
    raw_df: pd.DataFrame,
    forecast_days: int = 30,
    lookback_days: int = 60,
) -> dict:
    prices = raw_df["close"].values[-lookback_days:]
    S0 = float(prices[-1])

    # ── Estimate GBM parameters from daily log returns ───────────────
    log_ret = np.diff(np.log(prices))
    mu_daily = float(np.mean(log_ret))  # drift of log-returns
    sigma_daily = float(np.std(log_ret, ddof=1))  # daily volatility
    if not np.isfinite(sigma_daily) or sigma_daily < 1e-9:
        sigma_daily = 1e-9

    # ── Horizon grid (in trading days) ───────────────────────────────
    t = np.arange(1, forecast_days + 1, dtype=np.float64)

    # Expected price path under GBM: E[S_t] = S0 * exp(mu * t)
    # (using the log-return mean as drift; mu already includes -0.5 sigma^2
    #  implicitly because it is the mean of realised log returns).
    log_mean = np.log(S0) + mu_daily * t
    log_sd = sigma_daily * np.sqrt(t)  # sqrt-time uncertainty growth

    # Central projection = median of the lognormal (exp of the log-mean path).
    forecast_vals = np.exp(log_mean)

    # Lognormal quantile bands (multiplicative, always positive, fan out).
    upper_1 = np.exp(log_mean + 1.0 * log_sd)
    lower_1 = np.exp(log_mean - 1.0 * log_sd)
    upper_2 = np.exp(log_mean + 2.0 * log_sd)
    lower_2 = np.exp(log_mean - 2.0 * log_sd)

    # ── Future trading dates ─────────────────────────────────────────
    last_date = raw_df.index[-1]
    future_dates = pd.bdate_range(
        start=last_date + pd.Timedelta(days=1),
        periods=forecast_days,
    )

    # ── Trend direction from annualised drift ────────────────────────
    mu_annual_pct = mu_daily * 252 * 100
    if mu_daily * 100 > 0.05:
        direction = "bullish"
    elif mu_daily * 100 < -0.05:
        direction = "bearish"
    else:
        direction = "neutral"

    return {
        "forecast_dates": future_dates,
        "forecast_prices": forecast_vals,
        "upper_1sigma": upper_1,
        "lower_1sigma": lower_1,
        "upper_2sigma": upper_2,
        "lower_2sigma": lower_2,
        "last_price": S0,
        "trend_direction": direction,
        "trend_strength": round(mu_daily * S0, 4),  # ~$/day drift at S0
        "slope_pct_day": round(mu_daily * 100, 4),  # drift %/day
        "mu_annual_pct": round(mu_annual_pct, 2),  # annualised drift
        "sigma_annual_pct": round(
            sigma_daily * np.sqrt(252) * 100, 2
        ),  # annualised vol
        "model": "GBM (Black-Scholes price process)",
        "lookback_days": lookback_days,
    }


#  AGENT PORTFOLIO FORECAST
def _engineer_features_for_forecast(df: pd.DataFrame) -> pd.DataFrame:
    import preprocess as pp

    df = pp.handle_missing(df.copy())
    for fn in [
        pp.add_sma,
        pp.add_ema,
        pp.add_rsi,
        pp.add_macd,
        pp.add_bollinger_bands,
        pp.add_obv,
        pp.add_returns,
        pp.add_adx,
        pp.add_long_momentum,
        pp.add_roc,
        pp.add_realised_vol,
    ]:
        df = fn(df)
    df.dropna(inplace=True)

    feature_cols = list(df.columns)
    scaler = MinMaxScaler(feature_range=(0, 1))
    df[feature_cols] = scaler.fit_transform(df[feature_cols])
    return df


def _infer_lookback(model, n_features: int) -> int:
    expected = model.observation_space.shape[0]
    for n_scalars in [9, 5, 2]:
        market = expected - n_scalars
        if market > 0 and market % n_features == 0:
            return market // n_features
    return 30


def forecast_agent_portfolio(
    ticker: str,
    raw_df: pd.DataFrame,
    forecast_trend: dict,
    model_dir: str = "models",
    initial_capital: float = 10_000.0,
    transaction_cost: float = 0.001,
    n_simulations: int = 3,
) -> dict:
    if not SB3_AVAILABLE or not ENV_AVAILABLE:
        return {"model_loaded": False, "sim_values": [], "forecast_dates": []}

    # Try best_model first (current retrained discrete-action models), then
    # legacy tuned. Loading a stale *_tuned model would mismatch the obs shape.
    model_path = None
    import os

    for name in ["best_model", "best_model_tuned"]:
        path = f"{model_dir}/{ticker}/{name}.zip"
        if os.path.exists(path):
            model_path = path
            break

    if model_path is None:
        return {
            "model_loaded": False,
            "sim_values": [],
            "forecast_dates": [],
            "error": f"No model found in {model_dir}/{ticker}/",
        }

    model = PPO.load(model_path)

    forecast_dates = forecast_trend["forecast_dates"]
    forecast_prices = forecast_trend["forecast_prices"]
    sigma = forecast_trend["upper_1sigma"] - forecast_prices

    all_sim_values = []
    _last_shape = None

    for sim_idx in range(n_simulations):
        # Add small Gaussian noise to create scenario variation
        if sim_idx == 0:
            noise = np.zeros(len(forecast_prices))  # central
        elif sim_idx == 1:
            noise = 0.5 * sigma  # optimistic
        else:
            noise = -0.5 * sigma  # pessimistic

        sim_prices = np.maximum(forecast_prices + noise, 0.01)

        # Build a synthetic DataFrame extending the real history
        # with the forecasted prices so the agent has a proper
        # lookback window when it starts predicting
        extended_raw = raw_df.copy()

        for i, (date, price) in enumerate(zip(forecast_dates, sim_prices)):
            new_row = pd.DataFrame(
                {
                    "open": [price * 0.999],
                    "high": [price * 1.005],
                    "low": [price * 0.995],
                    "close": [price],
                    "volume": [extended_raw["volume"].mean()],
                },
                index=[date],
            )
            extended_raw = pd.concat([extended_raw, new_row])

        # Engineer features on the full extended series
        extended_feat = _engineer_features_for_forecast(extended_raw)

        # Trim to just the forecast period rows
        n_features = len(extended_feat.columns)
        lookback_window = _infer_lookback(model, n_features)

        # Only keep the last (lookback + forecast_days) rows
        forecast_len = len(forecast_dates)
        needed = lookback_window + forecast_len
        extended_feat = extended_feat.iloc[-needed:].reset_index(drop=True)

        # Build env on the forecast slice
        env = TradingEnv(
            df=extended_feat,
            initial_capital=initial_capital,
            base_cost_pct=transaction_cost,
            market_impact_pct=0.0005,
            overtrading_threshold=20,
            overtrading_penalty_pct=0.0001,
            lookback_window=lookback_window,
        )

        obs, _ = env.reset()

        # Check shape compatibility
        if obs.shape[0] != model.observation_space.shape[0]:
            env.close()
            _last_shape = (obs.shape[0], model.observation_space.shape[0])
            continue

        # Step the agent and record the direction it executes each day. The
        # env trades on NORMALISED prices, so its internal portfolio_value is
        # not real dollars and must NOT be rescaled. Instead we replay the agent's executed
        # directions on the REAL forecast prices to get a faithful dollar P&L.
        directions = []
        done = truncated = False
        while not done and not truncated:
            action, _ = model.predict(obs, deterministic=(sim_idx == 0))
            obs, _, done, truncated, info = env.step(action)
            directions.append(int(info["direction"]))

        env.close()

        # Replay on real prices: sim_prices are the GBM forecast in dollars.
        # Align directions to the forecast horizon (one decision per day).
        cash = float(initial_capital)
        shares = 0.0
        dollar_values = []
        n_steps = min(len(directions), len(sim_prices))
        for k in range(n_steps):
            price = float(sim_prices[k])
            d = directions[k]
            if d == 1 and cash > 0 and price > 1e-8:  # buy with all cash
                buy_price = price * (1.0 + 0.0005)
                qty = cash / buy_price
                cash -= qty * buy_price * (1.0 + transaction_cost)
                shares += qty
            elif d == 2 and shares > 0:  # sell all
                sell_price = price * (1.0 - 0.0005)
                cash += shares * sell_price * (1.0 - transaction_cost)
                shares = 0.0
            dollar_values.append(cash + shares * price)

        # Pad/trim to the forecast horizon
        dollar_values = dollar_values[:forecast_len]
        while len(dollar_values) < forecast_len:
            dollar_values.append(
                dollar_values[-1] if dollar_values else initial_capital
            )

        all_sim_values.append(dollar_values)

    result = {
        "model_loaded": True,
        "sim_values": all_sim_values,
        "forecast_dates": forecast_dates,
        "base_projection": all_sim_values[0] if all_sim_values else [],
    }
    if not all_sim_values and _last_shape is not None:
        result["error"] = (
            f"Observation shape mismatch: env produced {_last_shape[0]}, "
            f"model expects {_last_shape[1]}. The model was trained with a "
            f"different feature/lookback config — retrain or check features."
        )
    return result


#  PLOT — PRICE FORECAST
def plot_price_forecast(
    ticker: str,
    raw_df: pd.DataFrame,
    forecast: dict,
    theme: dict,
    history_days: int = 90,
) -> "matplotlib.figure.Figure":
    acc = theme["accent"]
    grn = theme["green"]
    red = theme["red"]
    bg = theme["bg"]
    panel = theme["panel"]
    bdr = theme["border"]
    muted = theme["muted"]
    txt = theme["text"]

    # Historical slice
    hist_prices = raw_df["close"].values[-history_days:]
    hist_dates = raw_df.index[-history_days:]
    returns = np.diff(hist_prices) / hist_prices[:-1] * 100

    fig = plt.figure(figsize=(13, 8), facecolor=bg)
    gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1], hspace=0.1, figure=fig)
    ax1 = fig.add_subplot(gs[0])
    ax2 = fig.add_subplot(gs[1])

    direction_color = (
        grn
        if forecast["trend_direction"] == "bullish"
        else (red if forecast["trend_direction"] == "bearish" else muted)
    )
    direction_label = forecast["trend_direction"].upper()

    fig.suptitle(
        f"  {ticker}  ·  {history_days}-Day History + {len(forecast['forecast_dates'])}-Day Forecast"
        f"  [{direction_label}  {forecast['slope_pct_day']:+.3f}%/day]",
        fontsize=13,
        fontweight="bold",
        color=direction_color,
        x=0.02,
        ha="left",
        y=0.98,
    )

    ax1.set_facecolor(panel)

    # Historical price
    ax1.plot(
        hist_dates,
        hist_prices,
        color=txt,
        linewidth=1.6,
        label="Historical Close",
        zorder=3,
    )

    # Confidence bands
    fd = forecast["forecast_dates"]
    fp = forecast["forecast_prices"]

    ax1.fill_between(
        fd,
        forecast["lower_2sigma"],
        forecast["upper_2sigma"],
        alpha=0.10,
        color=direction_color,
        label="±2σ band",
    )
    ax1.fill_between(
        fd,
        forecast["lower_1sigma"],
        forecast["upper_1sigma"],
        alpha=0.20,
        color=direction_color,
        label="±1σ band",
    )

    # Central forecast line
    ax1.plot(
        fd,
        fp,
        color=direction_color,
        linewidth=2.0,
        linestyle="--",
        label=f"Trend forecast",
        zorder=4,
    )

    # Vertical separator at today
    ax1.axvline(x=hist_dates[-1], color=muted, linewidth=1.0, linestyle=":", alpha=0.7)
    ax1.text(
        hist_dates[-1],
        ax1.get_ylim()[0]
        if ax1.get_ylim()[0] != ax1.get_ylim()[1]
        else hist_prices.min(),
        "  Today",
        color=muted,
        fontsize=8,
        va="bottom",
    )

    ax1.set_ylabel("Price ($)", color=muted, fontsize=10)
    ax1.yaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"${x:,.2f}"))
    ax1.legend(loc="upper left", fontsize=9, framealpha=0.9)
    ax1.grid(True, linestyle="--", alpha=0.35)
    plt.setp(ax1.get_xticklabels(), visible=False)

    # ── Return distribution ────────────────────────────────────
    ax2.set_facecolor(panel)

    pos_returns = returns[returns >= 0]
    neg_returns = returns[returns < 0]

    ax2.hist(pos_returns, bins=30, color=grn, alpha=0.7, label="Positive")
    ax2.hist(neg_returns, bins=30, color=red, alpha=0.7, label="Negative")
    ax2.axvline(x=0, color=muted, linewidth=0.8)
    ax2.axvline(
        x=float(np.mean(returns)),
        color=direction_color,
        linewidth=1.5,
        linestyle="--",
        label=f"Mean {np.mean(returns):+.2f}%",
    )

    ax2.set_ylabel("Frequency", color=muted, fontsize=9)
    ax2.set_xlabel("Daily Return (%)", color=muted, fontsize=9)
    ax2.set_title("Daily Return Distribution", color=txt, fontsize=10, pad=6)
    ax2.legend(fontsize=8, framealpha=0.85)
    ax2.grid(True, linestyle="--", alpha=0.3)

    for ax in [ax1, ax2]:
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.spines["bottom"].set_edgecolor(bdr)
        ax.spines["left"].set_edgecolor(bdr)

    plt.tight_layout(rect=(0, 0, 1, 0.96))
    return fig


#  PLOT — AGENT PORTFOLIO FORECAST
def plot_portfolio_forecast(
    ticker: str,
    forecast_trend: dict,
    portfolio_forecast: dict,
    initial_capital: float,
    theme: dict,
) -> "matplotlib.figure.Figure":
    """
    Plot the three simulated portfolio trajectories:
      Central (no noise), Optimistic (+0.5σ), Pessimistic (−0.5σ)
    alongside the trend-implied Buy & Hold baseline.
    """

    acc = theme["accent"]
    grn = theme["green"]
    red = theme["red"]
    amb = "#f59e0b"
    bg = theme["bg"]
    panel = theme["panel"]
    bdr = theme["border"]
    muted = theme["muted"]
    txt = theme["text"]

    fig, ax = plt.subplots(figsize=(13, 5), facecolor=bg)
    ax.set_facecolor(panel)

    dates = forecast_trend["forecast_dates"]

    if not portfolio_forecast["model_loaded"] or not portfolio_forecast["sim_values"]:
        ax.text(
            0.5,
            0.5,
            f"No trained model found for {ticker}.\nPortfolio forecast unavailable.",
            ha="center",
            va="center",
            color=muted,
            fontsize=12,
            transform=ax.transAxes,
        )
        ax.axis("off")
        return fig

    sim_values = portfolio_forecast["sim_values"]
    labels = ["Central (base case)", "Optimistic (+0.5σ)", "Pessimistic (−0.5σ)"]
    colors = [acc, grn, red]
    styles = ["-", "--", "-."]
    alphas = [1.0, 0.75, 0.75]

    for i, (vals, label, color, ls, alpha) in enumerate(
        zip(sim_values, labels, colors, styles, alphas)
    ):
        min_len = min(len(dates), len(vals))
        ret = (vals[min_len - 1] - initial_capital) / initial_capital * 100
        ax.plot(
            dates[:min_len],
            vals[:min_len],
            label=f"{label}  ({ret:+.1f}%)",
            color=color,
            linewidth=2.0 if i == 0 else 1.4,
            linestyle=ls,
            alpha=alpha,
        )

    # Buy & Hold baseline using trend forecast prices
    fp = forecast_trend["forecast_prices"]
    lp = forecast_trend["last_price"]
    if lp > 1e-8:
        bnh_shares = int(initial_capital // lp)
        bnh_vals = [initial_capital + bnh_shares * (p - lp) for p in fp]
    else:
        bnh_vals = [initial_capital] * len(fp)
    bnh_ret = (bnh_vals[-1] - initial_capital) / initial_capital * 100
    min_len = min(len(dates), len(bnh_vals))
    ax.plot(
        dates[:min_len],
        bnh_vals[:min_len],
        label=f"Buy & Hold (trend)  ({bnh_ret:+.1f}%)",
        color=amb,
        linewidth=1.4,
        linestyle=":",
        alpha=0.85,
    )

    ax.axhline(y=initial_capital, color=muted, linewidth=0.8, linestyle=":", alpha=0.5)

    ax.set_title(
        f"{ticker}  ·  Agent Portfolio Forecast — Next {len(dates)} Trading Days",
        color=acc,
        fontsize=12,
        loc="left",
    )
    ax.set_ylabel("Projected Portfolio Value ($)", color=muted, fontsize=10)
    ax.set_xlabel("Date", color=muted, fontsize=10)
    ax.yaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"${x:,.0f}"))
    ax.legend(loc="upper left", fontsize=9, framealpha=0.9)
    ax.grid(True, linestyle="--", alpha=0.35)

    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["bottom"].set_edgecolor(bdr)
    ax.spines["left"].set_edgecolor(bdr)

    # Disclaimer
    ax.text(
        0.99,
        0.02,
        "Forecast based on Geometric Brownian Motion (Black-Scholes). Not financial advice.",
        transform=ax.transAxes,
        ha="right",
        va="bottom",
        fontsize=8,
        color=muted,
        style="italic",
    )

    plt.tight_layout()
    return fig


def build_forecast_summary_html(
    ticker: str,
    forecast: dict,
    portfolio_forecast: dict,
    initial_capital: float,
    theme: dict,
) -> str:
    """
    Small HTML card summarising the forecast in numbers.
    Sits above the forecast charts in the dashboard.
    """

    acc = theme["accent"]
    grn = theme["green"]
    red = theme["red"]
    txt = theme["text"]
    muted = theme["muted"]
    panel = theme["panel"]
    bdr = theme["border"]

    direction = forecast["trend_direction"]
    color_map = {"bullish": grn, "bearish": red, "neutral": muted}
    dir_color = color_map.get(direction, muted)
    dir_emoji = {"bullish": "↑", "bearish": "↓", "neutral": "→"}.get(direction, "→")

    last_price = forecast["last_price"]
    end_price = forecast["forecast_prices"][-1]
    price_change = (end_price - last_price) / last_price * 100

    proj_end = ""
    if portfolio_forecast["model_loaded"] and portfolio_forecast["sim_values"]:
        vals = portfolio_forecast["sim_values"][0]
        proj_end_val = vals[-1] if vals else initial_capital
        proj_ret = (proj_end_val - initial_capital) / initial_capital * 100
        proj_end = (
            f'<div style="text-align:center;padding:12px 16px;'
            f'background:#eff6ff;border-radius:8px;border:1px solid #bfdbfe">'
            f'<div style="font-size:10px;color:{muted};letter-spacing:.08em;'
            f'text-transform:uppercase;margin-bottom:4px">Agent Projected Return</div>'
            f'<div style="font-size:20px;font-weight:700;color:{acc}">'
            f"{proj_ret:+.2f}%</div></div>"
        )

    n_days = len(forecast["forecast_dates"])

    html = f"""
<div style="
    font-family:'IBM Plex Mono',monospace;
    background:{panel};
    border:1px solid {bdr};
    border-radius:12px;
    padding:18px 22px;
    margin-bottom:12px;
">
    <div style="display:flex;align-items:center;gap:10px;margin-bottom:14px">
        <span style="
            background:{dir_color};color:#fff;
            font-size:10px;font-weight:600;
            letter-spacing:.1em;text-transform:uppercase;
            padding:3px 10px;border-radius:6px;
        ">{dir_emoji} {direction}</span>
        <span style="color:{muted};font-size:11px">
            {ticker} · {n_days}-Day Forecast
        </span>
        <span style="color:{muted};font-size:10px;margin-left:auto;font-style:italic">
            GBM · Black-Scholes price process · Not financial advice
        </span>
    </div>

    <div style="display:grid;grid-template-columns:repeat(4,1fr);gap:12px">

        <div style="text-align:center;padding:12px 16px;
                    background:#f8fafc;border-radius:8px;border:1px solid {bdr}">
            <div style="font-size:10px;color:{muted};letter-spacing:.08em;
                        text-transform:uppercase;margin-bottom:4px">Last Close</div>
            <div style="font-size:20px;font-weight:700;color:{txt}">${last_price:,.2f}</div>
        </div>

        <div style="text-align:center;padding:12px 16px;
                    background:#f8fafc;border-radius:8px;border:1px solid {bdr}">
            <div style="font-size:10px;color:{muted};letter-spacing:.08em;
                        text-transform:uppercase;margin-bottom:4px">Projected Price</div>
            <div style="font-size:20px;font-weight:700;color:{dir_color}">${end_price:,.2f}</div>
        </div>

        <div style="text-align:center;padding:12px 16px;
                    background:#f8fafc;border-radius:8px;border:1px solid {bdr}">
            <div style="font-size:10px;color:{muted};letter-spacing:.08em;
                        text-transform:uppercase;margin-bottom:4px">Price Change</div>
            <div style="font-size:20px;font-weight:700;color:{dir_color}">{price_change:+.2f}%</div>
        </div>

        {proj_end}

    </div>
</div>
"""
    return html