main.py

import os

import numpy as np
from dotenv import load_dotenv
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Confirm, IntPrompt, Prompt
from rich.table import Table

from src.analysis import (
    calculate_max_drawdown,
    calculate_sharpe_ratio,
    create_sparkline,
)
from src.backtest import run_backtest_logic

load_dotenv()
console = Console()

POPULAR_TICKERS = [
    ("Apple", "AAPL"),
    ("Microsoft", "MSFT"),
    ("Google", "GOOGL"),
    ("Amazon", "AMZN"),
    ("NVIDIA", "NVDA"),
    ("Meta", "META"),
    ("Tesla", "TSLA"),
    ("Netflix", "NFLX"),
    ("AMD", "AMD"),
    ("Intel", "INTC"),
    ("Bitcoin", "BTC-USD"),
    ("Ethereum", "ETH-USD"),
    ("Gold", "GC=F"),
    ("Oil", "CL=F"),
    ("SPY", "SPY"),
    ("QQQ", "QQQ"),
]

QUICK_PRESETS = {
    "1": {"name": "Ultra-Fast (5 days, every 2nd day)", "days": 5, "skip": 2},
    "2": {"name": "Quick (10 days, daily)", "days": 10, "skip": 1},
    "3": {"name": "Standard (30 days, daily)", "days": 30, "skip": 1},
    "4": {"name": "Deep (60 days, daily)", "days": 60, "skip": 1},
    "5": {"name": "Custom", "days": None, "skip": None},
}


def display_welcome():
    console.print(
        Panel.fit(
            "[bold green]🚀 Quant Terminal[/bold green]\n"
            "[dim]AI-Powered Backtesting with Quick-Run Presets[/dim]",
            border_style="cyan",
        )
    )


def display_available_stocks():
    table = Table(title="Available Stocks", show_header=True, header_style="bold cyan")
    table.add_column("#", justify="right", style="dim")
    table.add_column("Name", style="dim")
    table.add_column("Ticker")

    for i, (name, ticker) in enumerate(POPULAR_TICKERS, 1):
        table.add_row(str(i), name, ticker)

    console.print(table)


def display_presets():
    table = Table(
        title="Quick-Run Presets", show_header=True, header_style="bold green"
    )
    table.add_column("#", justify="right", style="dim")
    table.add_column("Preset", style="bold")
    table.add_column("Duration", style="dim")
    table.add_column("Speed", style="dim")

    for key, preset in QUICK_PRESETS.items():
        if preset["days"]:
            days_str = f"{preset['days']} days"
            skip_str = f"every {preset['skip']}d" if preset["skip"] > 1 else "daily"
        else:
            days_str = "Custom"
            skip_str = "Custom"
        table.add_row(key, preset["name"], days_str, skip_str)

    console.print(table)


def get_user_input():
    display_available_stocks()
    console.print(
        "\n[bold yellow]Choose a stock by number or enter a custom ticker:[/bold yellow]"
    )
    ticker = Prompt.ask("[bold yellow]Ticker[/bold yellow]", default="AAPL").upper()

    if ticker.isdigit() and 1 <= int(ticker) <= len(POPULAR_TICKERS):
        idx = int(ticker) - 1
        name = POPULAR_TICKERS[idx][0]
        ticker = POPULAR_TICKERS[idx][1]
        console.print(f"[bold green]Selected: {name} ({ticker})[/bold green]")

    display_presets()
    console.print(
        "\n[bold yellow]Choose a preset (1-5) or enter custom settings:[/bold yellow]"
    )
    preset_choice = Prompt.ask("[bold yellow]Preset[/bold yellow]", default="2")

    if preset_choice in QUICK_PRESETS:
        preset = QUICK_PRESETS[preset_choice]
        if preset["days"] is None:
            days = IntPrompt.ask("Enter backtest duration (days)", default=30)
            skip_days = IntPrompt.ask(
                "Trade every Nth day (1=daily, 2=every 2 days)", default=1
            )
        else:
            days = preset["days"]
            skip_days = preset["skip"]
            console.print(f"[bold green]Using preset: {preset['name']}[/bold green]")
    else:
        days = IntPrompt.ask("Enter backtest duration (days)", default=30)
        skip_days = IntPrompt.ask(
            "Trade every Nth day (1=daily, 2=every 2 days)", default=1
        )

    return ticker, days, skip_days


def display_report(results):
    if not results:
        return

    portfolio = results["portfolio"]
    data = results["data"]
    ticker = results["ticker"]
    completed_trades = results["completed_trades"]
    held_periods = results["held_periods"]
    decisions_made = results.get("decisions_made", 0)

    final_price = data["Close"].iloc[-1]
    total_return = portfolio.get_total_return()
    start_price = data["Close"].iloc[20] if len(data) > 20 else data["Close"].iloc[0]
    buy_hold_return = ((final_price - start_price) / start_price) * 100

    perf_table = Table(
        title=f"Performance Report: {ticker}",
        show_header=True,
        header_style="bold magenta",
    )
    perf_table.add_column("Metric", style="dim")
    perf_table.add_column("Value", justify="right")

    perf_table.add_row("Initial Balance", f"${portfolio.initial_balance:,.2f}")
    perf_table.add_row("Final Balance", f"${portfolio.balance:,.2f}")
    perf_table.add_row("Portfolio Value", f"${portfolio.get_portfolio_value():,.2f}")

    ret_style = "green" if total_return >= 0 else "red"
    perf_table.add_row(
        "Total Return", f"[{ret_style}]{total_return:.2f}%[/{ret_style}]"
    )
    perf_table.add_row(
        "Alpha (vs Buy & Hold)", f"{total_return - buy_hold_return:.2f}%"
    )

    max_dd = calculate_max_drawdown(portfolio.equity_curve)
    returns = portfolio.get_cumulative_returns()
    sharpe = calculate_sharpe_ratio(returns)

    risk_table = Table(title="Risk Metrics", show_header=True, header_style="bold red")
    risk_table.add_column("Metric", style="dim")
    risk_table.add_column("Value", justify="right")
    risk_table.add_row("Max Drawdown", f"-{max_dd:.2f}%")
    risk_table.add_row("Sharpe Ratio", f"{sharpe:.4f}")
    risk_table.add_row(
        "Volatility", f"{np.std(returns):.4f}%" if len(returns) > 0 else "N/A"
    )

    stats_table = Table(
        title="Trade Statistics", show_header=True, header_style="bold blue"
    )
    stats_table.add_column("Stat", style="dim")
    stats_table.add_column("Value", justify="right")
    stats_table.add_row("Decisions Made", str(decisions_made))
    stats_table.add_row("Completed Trades", str(completed_trades))
    stats_table.add_row("Hold Periods", str(held_periods))
    if decisions_made > 0:
        trade_rate = (completed_trades / decisions_made) * 100
        stats_table.add_row("Trade Rate", f"{trade_rate:.1f}%")

    prices = data["Close"].tolist()
    spark = create_sparkline(prices[-20:])

    console.print("\n")
    console.print(perf_table)
    console.print(risk_table)
    console.print(stats_table)
    console.print(f"\n[bold cyan]Price Trend (Last 20 steps):[/bold cyan] {spark}")

    if portfolio.trades:
        trade_table = Table(
            title="Trade History", show_header=True, header_style="bold blue"
        )
        trade_table.add_column("Date")
        trade_table.add_column("Action")
        trade_table.add_column("Price")
        trade_table.add_column("Balance")
        trade_table.add_column("Reasoning", ratio=1)

        for trade in portfolio.trades[-10:]:
            action_color = (
                "green"
                if trade["action"] == "BUY"
                else "red"
                if trade["action"] == "SELL"
                else "yellow"
            )
            trade_table.add_row(
                trade["date"],
                f"[{action_color}]{trade['action']}[/{action_color}]",
                f"${trade['price']:.2f}",
                f"${trade['balance']:,.2f}",
                trade["reasoning"][:60],
            )
        console.print("\n")
        console.print(trade_table)


def main():
    display_welcome()
    ticker, days, skip_days = get_user_input()

    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        console.print("[bold red]Error: GEMINI_API_KEY not found in .env[/bold red]")
        return

    console.print(
        f"\n[bold cyan]Running backtest for {ticker} ({days} days, skip={skip_days})...[/bold cyan]\n"
    )

    def log_callback(message):
        console.print(message)

    results = run_backtest_logic(
        ticker=ticker,
        lookback_days=days,
        api_key=api_key,
        log_callback=log_callback,
        skip_days=skip_days,
    )

    console.print("\n")
    if results:
        display_report(results)
    else:
        console.print("[bold red]Backtest failed or returned no data.[/bold red]")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        console.print("\n[yellow]Exiting...[/yellow]")