import os
import pickle
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.ticker as mticker
from typing import Optional

RESULTS_DIR   = "results"
PLOTS_DIR     = os.path.join(RESULTS_DIR, "plots")
VALID_TICKERS = ["AAPL", "MSFT", "AMZN", "GOOGL", "NVDA"]
os.makedirs(PLOTS_DIR, exist_ok=True)


def load_results() -> tuple:
    backtest_path = os.path.join(RESULTS_DIR, "backtest_results.pkl")
    baseline_path = os.path.join(RESULTS_DIR, "baseline_results.pkl")

    if not os.path.exists(backtest_path):
        raise FileNotFoundError(f"Run evaluate.py first — {backtest_path} missing.")
    if not os.path.exists(baseline_path):
        raise FileNotFoundError(f"Run baselines.py first — {baseline_path} missing.")

    with open(backtest_path, "rb") as f:
        backtest_results = pickle.load(f)
    with open(baseline_path, "rb") as f:
        baseline_results = pickle.load(f)

    print(f"Loaded backtest results  : {list(backtest_results.keys())}")
    print(f"Loaded baseline results  : {list(baseline_results.keys())}")
    return backtest_results, baseline_results


def calc_cumulative_return(values: np.ndarray, initial: float) -> float:
    return (values[-1] - initial) / initial


def calc_annualised_return(values: np.ndarray, trading_days: int = 252) -> float:
    n_years = len(values) / trading_days
    if n_years <= 0 or values[0] <= 0:
        return 0.0
    return float((values[-1] / values[0]) ** (1 / n_years) - 1)


def calc_annualised_volatility(values: np.ndarray, trading_days: int = 252) -> float:
    returns = np.diff(values) / values[:-1]
    return float(np.std(returns) * np.sqrt(trading_days))


def calc_sharpe_ratio(
    values: np.ndarray,
    risk_free_rate: float = 0.04,
    trading_days: int = 252,
) -> float:
    returns  = np.diff(values) / values[:-1]
    daily_rf = risk_free_rate / trading_days
    excess   = returns - daily_rf
    if np.std(excess) == 0:
        return 0.0
    return float(np.mean(excess) / np.std(excess) * np.sqrt(trading_days))


def calc_max_drawdown(values: np.ndarray) -> float:
    peak      = np.maximum.accumulate(values)
    drawdowns = (values - peak) / peak
    return float(np.min(drawdowns))


def calc_sortino_ratio(
    values: np.ndarray,
    risk_free_rate: float = 0.04,
    trading_days: int = 252,
) -> float:
    returns  = np.diff(values) / values[:-1]
    daily_rf = risk_free_rate / trading_days
    excess   = returns - daily_rf
    downside = excess[excess < 0]
    if len(downside) == 0 or np.std(downside) == 0:
        return 0.0
    downside_std = float(np.std(downside) * np.sqrt(trading_days))
    return float(np.mean(excess) * trading_days / downside_std)


def calc_calmar_ratio(values: np.ndarray, trading_days: int = 252) -> float:
    ann_return = calc_annualised_return(values, trading_days)
    max_dd     = calc_max_drawdown(values)
    if max_dd == 0:
        return 0.0
    return float(ann_return / abs(max_dd))


def compute_all_metrics(
    values: np.ndarray,
    initial_capital: float,
    total_trades: Optional[int] = None,
    risk_free_rate: float = 0.04,
    trading_days: int = 252,
) -> dict:
    values = np.array(values, dtype=np.float64)
    return {
        "Final Value ($)":       round(float(values[-1]), 2),
        "Cumulative Return":     round(calc_cumulative_return(values, initial_capital), 4),
        "Annualised Return":     round(calc_annualised_return(values, trading_days), 4),
        "Annualised Volatility": round(calc_annualised_volatility(values, trading_days), 4),
        "Sharpe Ratio":          round(calc_sharpe_ratio(values, risk_free_rate, trading_days), 4),
        "Max Drawdown":          round(calc_max_drawdown(values), 4),
        "Sortino Ratio":         round(calc_sortino_ratio(values, risk_free_rate, trading_days), 4),
        "Calmar Ratio":          round(calc_calmar_ratio(values, trading_days), 4),
        "Total Trades":          total_trades if total_trades is not None else "N/A",
    }


def compute_drawdown_series(values: np.ndarray) -> np.ndarray:
    values = np.array(values, dtype=np.float64)
    peak   = np.maximum.accumulate(values)
    return (values - peak) / peak * 100


def _simulate_buy_and_hold(
    df: pd.DataFrame,
    initial_capital: float,
    transaction_cost: float = 0.001,
) -> tuple:
    cash   = initial_capital
    shares = 0
    values = []

    entry_price = float(df.iloc[0]["close"])
    if entry_price > 1e-8:
        shares  = int(cash // entry_price)
        cost    = shares * entry_price
        cash   -= cost * (1 + transaction_cost)

    for i in range(len(df)):
        price = float(df.iloc[i]["close"])
        values.append(cash + shares * price)

    exit_price  = float(df.iloc[-1]["close"])
    if shares > 0 and exit_price > 1e-8:
        final_cash  = shares * exit_price * (1 - transaction_cost) + cash
        values[-1]  = final_cash

    return values, 2


def _simulate_sma_crossover(
    df: pd.DataFrame,
    initial_capital: float,
    short_window: int = 20,
    long_window: int  = 50,
    transaction_cost: float = 0.001,
) -> tuple:
    cash   = initial_capital
    shares = 0
    values = []
    trades = 0
    pos    = "out"

    if "sma_20" in df.columns and "sma_50" in df.columns:
        sma_s = df["sma_20"].values
        sma_l = df["sma_50"].values
    else:
        sma_s = df["close"].rolling(short_window).mean().values
        sma_l = df["close"].rolling(long_window).mean().values

    close = df["close"].values

    for i in range(len(df)):
        price = float(close[i])

        if (i > 0
                and not np.isnan(sma_s[i]) and not np.isnan(sma_l[i])
                and not np.isnan(sma_s[i-1]) and not np.isnan(sma_l[i-1])):

            prev_above = sma_s[i-1] > sma_l[i-1]
            curr_above = sma_s[i]   > sma_l[i]

            if not prev_above and curr_above and pos == "out" and price > 1e-8:
                shares  = int(cash // price)
                cost    = shares * price
                cash   -= cost * (1 + transaction_cost)
                pos     = "in"
                trades += 1

            elif prev_above and not curr_above and pos == "in":
                proceeds = shares * price * (1 - transaction_cost)
                cash    += proceeds
                shares   = 0
                pos      = "out"
                trades  += 1

        values.append(cash + shares * price)

    if shares > 0:
        proceeds   = shares * float(close[-1]) * (1 - transaction_cost)
        cash      += proceeds
        values[-1] = cash
        trades    += 1

    return values, trades


def build_comparison_table(
    ticker: str,
    backtest_results: dict,
    baseline_results: dict,
    initial_capital: float = 10_000.0,
) -> pd.DataFrame:
    rows    = {}
    test_df = backtest_results[ticker]["test_df"]

    rl_values = backtest_results[ticker]["history_df"]["portfolio_values"].values
    rl_trades = backtest_results[ticker]["total_trades"]
    rows["RL Agent (PPO)"] = compute_all_metrics(rl_values, initial_capital, rl_trades)

    bnh_values, bnh_trades = _simulate_buy_and_hold(test_df, initial_capital)
    rows["Buy and Hold"]   = compute_all_metrics(bnh_values, initial_capital, bnh_trades)

    sma_values, sma_trades = _simulate_sma_crossover(test_df, initial_capital)
    rows["SMA Crossover"]  = compute_all_metrics(sma_values, initial_capital, sma_trades)

    return pd.DataFrame(rows)


def print_kpi_table(ticker: str, df: pd.DataFrame) -> None:
    FORMAT_MAP = {
        "Final Value ($)":       lambda v: f"${v:>12,.2f}",
        "Cumulative Return":     lambda v: f"{v*100:>11.2f}%",
        "Annualised Return":     lambda v: f"{v*100:>11.2f}%",
        "Annualised Volatility": lambda v: f"{v*100:>11.2f}%",
        "Sharpe Ratio":          lambda v: f"{v:>12.4f}",
        "Max Drawdown":          lambda v: f"{v*100:>11.2f}%",
        "Sortino Ratio":         lambda v: f"{v:>12.4f}",
        "Calmar Ratio":          lambda v: f"{v:>12.4f}",
        "Total Trades":          lambda v: f"{v:>12}",
    }
    LOWER_IS_BETTER = {"Annualised Volatility", "Max Drawdown"}
    col_w   = 18
    label_w = 26

    print(f"\n{'-'*75}")
    print(f"  KPI Comparison — {ticker}")
    print(f"{'-'*75}")
    print(f"{'Metric':<{label_w}}" + "".join(
        f"{col:>{col_w}}" for col in df.columns
    ))
    print(f"{'─'*75}")

    for metric in df.index:
        if metric not in FORMAT_MAP:
            continue
        row_values   = df.loc[metric]
        fmt          = FORMAT_MAP[metric]
        numeric_vals = {}

        for col in df.columns:
            val = row_values[col]
            if isinstance(val, (int, float)) and not np.isnan(float(val)):
                numeric_vals[col] = float(val)

        if numeric_vals:
            best_col = (min if metric in LOWER_IS_BETTER else max)(
                numeric_vals, key=numeric_vals.get
            )
        else:
            best_col = None

        row = f"  {metric:<{label_w-2}}"
        for col in df.columns:
            val = row_values[col]
            try:
                formatted = fmt(float(val))
                if col == best_col:
                    formatted = f"{formatted.strip() + ' ✓':>{col_w}}"
            except (ValueError, TypeError):
                formatted = f"{'N/A':>{col_w}}"
            row += formatted
        print(row)

    print(f"{'-'*75}")


def compute_drawdown_series(values: np.ndarray) -> np.ndarray:
    values = np.array(values, dtype=np.float64)
    peak   = np.maximum.accumulate(values)
    return (values - peak) / peak * 100


def plot_equity_curves(
    ticker: str,
    backtest_results: dict,
    initial_capital: float = 10_000.0,
    save: bool = True,
) -> None:
    if ticker not in backtest_results:
        print(f"  No backtest results for {ticker} — skipping plot.")
        return

    test_df = backtest_results[ticker]["test_df"]
    dates   = test_df.index

    rl_values              = backtest_results[ticker]["history_df"]["portfolio_values"].values
    bnh_values, _          = _simulate_buy_and_hold(test_df, initial_capital)
    sma_values, _          = _simulate_sma_crossover(test_df, initial_capital)
    bnh_values             = np.array(bnh_values)
    sma_values             = np.array(sma_values)

    min_len    = min(len(rl_values), len(bnh_values), len(sma_values), len(dates))
    dates      = dates[-min_len:]
    rl_values  = rl_values[-min_len:]
    bnh_values = bnh_values[-min_len:]
    sma_values = sma_values[-min_len:]

    rl_dd  = compute_drawdown_series(rl_values)
    bnh_dd = compute_drawdown_series(bnh_values)
    sma_dd = compute_drawdown_series(sma_values)

    COLORS = {
        "RL Agent (PPO)": "#2563EB",
        "Buy and Hold":   "#16A34A",
        "SMA Crossover":  "#DC2626",
    }
    STYLES = {
        "RL Agent (PPO)": "-",
        "Buy and Hold":   "--",
        "SMA Crossover":  "-.",
    }

    fig = plt.figure(figsize=(14, 9))
    fig.suptitle(
        f"{ticker} — Strategy Comparison  "
        f"({dates[0].date()} → {dates[-1].date()})",
        fontsize=14, fontweight="bold", y=0.98,
    )
    gs     = gridspec.GridSpec(2, 1, height_ratios=[3, 1], hspace=0.08)
    ax_top = fig.add_subplot(gs[0])
    ax_bot = fig.add_subplot(gs[1], sharex=ax_top)

    for label, values in [
        ("RL Agent (PPO)", rl_values),
        ("Buy and Hold",   bnh_values),
        ("SMA Crossover",  sma_values),
    ]:
        final_return = (values[-1] - initial_capital) / initial_capital * 100
        full_label   = f"{label}  ({final_return:+.1f}%)"
        ax_top.plot(
            dates, values,
            label     = full_label,
            color     = COLORS[label],
            linestyle = STYLES[label],
            linewidth = 2.0 if label == "RL Agent (PPO)" else 1.5,
            alpha     = 1.0 if label == "RL Agent (PPO)" else 0.85,
            zorder    = 3 if label == "RL Agent (PPO)" else 2,
        )

    ax_top.axhline(
        y=initial_capital, color="gray",
        linestyle=":", linewidth=1.0, alpha=0.6, label="Initial Capital",
    )
    ax_top.set_ylabel("Portfolio Value ($)", fontsize=11)
    ax_top.yaxis.set_major_formatter(
        mticker.FuncFormatter(lambda x, _: f"${x:,.0f}")
    )
    ax_top.legend(loc="upper left", fontsize=10, framealpha=0.9)
    ax_top.grid(True, alpha=0.3, linestyle="--")
    ax_top.set_title("Portfolio Value Over Time", fontsize=11, pad=8)
    plt.setp(ax_top.get_xticklabels(), visible=False)

    for label, dd in [
        ("RL Agent (PPO)", rl_dd),
        ("Buy and Hold",   bnh_dd),
        ("SMA Crossover",  sma_dd),
    ]:
        ax_bot.fill_between(dates, dd, 0, alpha=0.25, color=COLORS[label])
        ax_bot.plot(
            dates, dd,
            color=COLORS[label], linestyle=STYLES[label], linewidth=1.5,
        )

    ax_bot.set_ylabel("Drawdown (%)", fontsize=11)
    ax_bot.set_xlabel("Date", fontsize=11)
    ax_bot.yaxis.set_major_formatter(
        mticker.FuncFormatter(lambda x, _: f"{x:.0f}%")
    )
    ax_bot.grid(True, alpha=0.3, linestyle="--")
    ax_bot.set_title("Drawdown from Peak", fontsize=11, pad=8)

    plt.tight_layout()

    if save:
        path = os.path.join(PLOTS_DIR, f"{ticker}_equity_curve.png")
        plt.savefig(path, dpi=150, bbox_inches="tight")
        print(f"  Plot saved → {path}")

    plt.show()
    plt.close()


def save_kpi_tables(all_tables: dict) -> None:
    for ticker, df in all_tables.items():
        path = os.path.join(RESULTS_DIR, f"{ticker}_kpi_table.csv")
        df.to_csv(path)
        print(f"  KPI table saved → {path}")


def print_cross_ticker_summary(all_tables: dict) -> None:
    print(f"\n{'+'*65}")
    print(f"  Project Summary — RL Agent Across All Tickers")
    print(f"{'+'*65}")
    print(f"  {'Ticker':<8} {'Cum. Return':>13} {'Sharpe':>10} "
          f"{'Sortino':>10} {'Max DD':>10}")
    print(f"{'─'*65}")

    for ticker, table in all_tables.items():
        col = "RL Agent (PPO)"
        if col not in table.columns:
            continue
        cum_ret = table.loc["Cumulative Return",     col]
        sharpe  = table.loc["Sharpe Ratio",          col]
        sortino = table.loc["Sortino Ratio",         col]
        max_dd  = table.loc["Max Drawdown",          col]
        print(
            f"  {ticker:<8} "
            f"{cum_ret*100:>12.2f}% "
            f"{sharpe:>10.4f} "
            f"{sortino:>10.4f} "
            f"{max_dd*100:>10.2f}%"
        )

    print(f"{'-'*65}")
    print(f"\n  Plots  → {PLOTS_DIR}/")
    print(f"  Tables → {RESULTS_DIR}/")


def run_analysis(
    tickers: list          = VALID_TICKERS,
    initial_capital: float = 10_000.0,
) -> dict:

    backtest_results, baseline_results = load_results()
    all_tables = {}

    for ticker in tickers:
        if ticker not in backtest_results:
            print(f"\n  Skipping {ticker} — no backtest results.")
            continue

        print(f"\nAnalysing {ticker}...")
        table = build_comparison_table(
            ticker           = ticker,
            backtest_results = backtest_results,
            baseline_results = baseline_results,
            initial_capital  = initial_capital,
        )
        all_tables[ticker] = table
        print_kpi_table(ticker, table)
        plot_equity_curves(
            ticker           = ticker,
            backtest_results = backtest_results,
            initial_capital  = initial_capital,
            save             = True,
        )

    print("\nSaving KPI tables...")
    save_kpi_tables(all_tables)
    print_cross_ticker_summary(all_tables)

    return all_tables


if __name__ == "__main__":
    tables = run_analysis()
    print("\nAnalysis complete.")