%pip install --quiet pandas matplotlib seaborn numpy
# smtplib is built-in — nothing to install for email!
print("Libraries ready!")

# -- Standard library (built into Python) ------------------------------------
import smtplib                           # Sends emails via SMTP protocol
import ssl                               # Encrypted connection to Gmail
import time                              # For sleep/backoff between retries
import random                            # For jitter in backoff timing
import logging                           # Prints info/warning messages
from collections import deque            # List with max length (rolling window)
from datetime import datetime            # To timestamp each alert
from email.mime.text import MIMEText     # Plain text email body
from email.mime.multipart import MIMEMultipart  # Email with subject + body
from typing import List

# -- Third-party libraries (installed above) ----------------------------------
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.patches import Patch

# -- Logging setup ------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(message)s",
    datefmt="%H:%M:%S"
)
logger = logging.getLogger(__name__)
logger.info("All libraries imported successfully!")

# -- EDIT THESE TWO LINES -------------------------------------------------------
SENDER_EMAIL       = "arsalanbakhtiar@neurog.ai"    # Gmail account sending alerts
GMAIL_APP_PASSWORD = "pcev gxoo dtsz ifdb"     # 16-digit App Password from Google
# RECIPIENT_EMAIL    = "jahanzebahmed@neurog.ai"  # Who receives the alerts
RECIPIENT_EMAIL    = "razalens@gmail.com"  # Who receives the alerts
# -------------------------------------------------------------------------------

SMTP_CONFIG = {
    "host":     "smtp.gmail.com",
    "port":     465,           # Port 465 uses SSL (encrypted from the start)
    "sender":   SENDER_EMAIL,
    "password": GMAIL_APP_PASSWORD,
}

print(f"Sender   : {SENDER_EMAIL}")
print(f"Recipient: {RECIPIENT_EMAIL}")
print(f"SMTP     : {SMTP_CONFIG['host']}:{SMTP_CONFIG['port']}")
print("Config loaded!")

def create_alert_state(notification_email: str, max_history: int = 50) -> dict:
    """
    Initialize the alert system state dictionary.

    Parameters
    ----------
    notification_email : str
        Email address that will receive all trade alerts.
    max_history : int, optional
        Maximum number of past alerts to keep in memory, defaults to 50.

    Returns
    -------
    dict
        Initialized state with counters, history deque, and config.

    Examples
    --------
    >>> state = create_alert_state("trader@gmail.com")
    >>> state['sent_count']
    0
    """
    logger.info(f"Creating alert state for recipient: {notification_email}")
    state = {
        "notification_email": notification_email,
        "sent_count":         0,                       # Emails successfully sent
        "failed_count":       0,                       # Emails that failed
        "history":            deque(maxlen=max_history),  # Rolling window of alerts
        "latencies":          [],                      # Time (s) each send took
        "created_at":         datetime.now().isoformat(),
    }
    logger.debug(f"State created with max_history={max_history}")
    return state


# Quick test
state = create_alert_state(RECIPIENT_EMAIL)
print("State keys:", list(state.keys()))
print("Recipient :", state['notification_email'])

def format_trade_alert(
    state: dict,
    symbol: str,
    action: str,
    price: float,
    quantity: int
) -> dict:
    """
    Format raw trade data into an email-ready alert payload.

    Parameters
    ----------
    state : dict
        Current alert system state.
    symbol : str
        Stock ticker symbol, e.g. 'AAPL'.
    action : str
        Trade direction: 'BUY' or 'SELL'. Case-insensitive.
    price : float
        Execution price per share.
    quantity : int
        Number of shares traded.

    Returns
    -------
    dict
        Alert payload with subject, body, and metadata.

    Examples
    --------
    >>> alert = format_trade_alert(state, 'AAPL', 'BUY', 189.50, 10)
    >>> alert['subject']
    'BUY Alert: AAPL @ $189.50'
    """
    action      = action.upper()               # Normalize to uppercase
    total_value = price * quantity
    timestamp   = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    emoji       = "BUY" if action == "BUY" else "SELL"

    subject = f"{emoji} Alert: {symbol} @ ${price:.2f}"

    body = (
        f"Trade Alert\n"
        f"{'='*40}\n"
        f"Action   : {action}\n"
        f"Symbol   : {symbol}\n"
        f"Price    : ${price:.2f}\n"
        f"Quantity : {quantity} shares\n"
        f"Value    : ${total_value:,.2f}\n"
        f"Time     : {timestamp}\n"
        f"{'='*40}\n"
        f"Sent to  : {state['notification_email']}\n"
    )

    alert = {
        "subject":   subject,
        "body":      body,
        "symbol":    symbol,
        "action":    action,
        "price":     price,
        "quantity":  quantity,
        "value":     total_value,
        "timestamp": timestamp,
        "status":    "pending",   # Updated to 'sent' or 'failed' later
    }

    logger.info(f"Formatted: {action} {quantity}x {symbol} @ ${price:.2f}")
    return alert


# Quick test
sample_alert = format_trade_alert(state, "AAPL", "BUY", 189.50, 10)
print(sample_alert['subject'])
print()
print(sample_alert['body'])

def send_email_alert(state: dict, alert: dict, smtp_config: dict) -> dict:
    """
    Send a trade alert email via Gmail SMTP with exponential backoff retry.

    Uses smtplib.SMTP_SSL for an encrypted connection on port 465.
    Retries up to 3 times with exponential backoff and random jitter.

    Parameters
    ----------
    state : dict
        Current alert system state. Updated with counts and history.
    alert : dict
        Formatted alert payload from format_trade_alert().
    smtp_config : dict
        Must contain keys: 'host', 'port', 'sender', 'password'.

    Returns
    -------
    dict
        Updated state with sent_count, failed_count, latencies, and history.

    Examples
    --------
    >>> state = send_email_alert(state, alert, SMTP_CONFIG)
    >>> state['sent_count']
    1
    """
    max_retries = 3
    base_delay  = 1.0   # Seconds — doubles with each retry

    # Build the MIME email object
    msg            = MIMEMultipart()
    msg["From"]    = smtp_config["sender"]
    msg["To"]      = state["notification_email"]
    msg["Subject"] = alert["subject"]
    msg.attach(MIMEText(alert["body"], "plain"))

    for attempt in range(1, max_retries + 1):
        try:
            logger.info(f"Sending email attempt {attempt}/{max_retries}: {alert['subject']}")
            start_time = time.time()

            # ssl.create_default_context() verifies Google's certificate automatically
            context = ssl.create_default_context()

            with smtplib.SMTP_SSL(smtp_config["host"], smtp_config["port"], context=context) as server:
                server.login(smtp_config["sender"], smtp_config["password"])
                server.sendmail(
                    smtp_config["sender"],
                    state["notification_email"],
                    msg.as_string()
                )

            elapsed              = time.time() - start_time
            state["latencies"].append(elapsed)
            state["sent_count"] += 1
            alert["status"]      = "sent"
            logger.info(f"Email sent in {elapsed:.2f}s | total sent: {state['sent_count']}")
            break   # Success — stop retrying

        except smtplib.SMTPAuthenticationError:
            # Wrong credentials — no point retrying
            logger.error("Authentication failed. Check your App Password and sender email.")
            state["failed_count"] += 1
            alert["status"]        = "failed"
            break

        except Exception as e:
            logger.warning(f"Attempt {attempt} failed: {e}")
            if attempt < max_retries:
                # Exponential backoff: 1s, 2s, 4s — plus random jitter (0 to 0.5s)
                delay = base_delay * (2 ** (attempt - 1)) + random.uniform(0, 0.5)
                logger.info(f"Retrying in {delay:.1f}s...")
                time.sleep(delay)
            else:
                logger.error(f"All {max_retries} attempts failed for: {alert['subject']}")
                state["failed_count"] += 1
                alert["status"]        = "failed"

    # Save to rolling history — deque auto-drops oldest if full
    state["history"].append(alert.copy())
    return state


print("send_email_alert() defined")
print("NOTE: Set SENDER_EMAIL and GMAIL_APP_PASSWORD above to send real emails.")

def track_alert_metrics(state: dict) -> dict:
    """
    Compute performance metrics from the current alert state.

    Parameters
    ----------
    state : dict
        Current alert system state with history and counters.

    Returns
    -------
    dict
        Metrics: total_alerts, sent_count, failed_count, success_rate,
        avg_latency_s, total_value_usd, buy_count, sell_count.

    Examples
    --------
    >>> metrics = track_alert_metrics(state)
    >>> metrics['success_rate']
    100.0
    """
    total        = state["sent_count"] + state["failed_count"]
    success_rate = (state["sent_count"] / total * 100) if total > 0 else 0.0

    latencies   = state["latencies"]
    avg_latency = sum(latencies) / len(latencies) if latencies else 0.0

    history_list = list(state["history"])
    total_value  = sum(a.get("value", 0) for a in history_list)
    buy_count    = sum(1 for a in history_list if a.get("action") == "BUY")
    sell_count   = sum(1 for a in history_list if a.get("action") == "SELL")

    metrics = {
        "total_alerts":    total,
        "sent_count":      state["sent_count"],
        "failed_count":    state["failed_count"],
        "success_rate":    round(success_rate, 1),
        "avg_latency_s":   round(avg_latency, 3),
        "total_value_usd": round(total_value, 2),
        "buy_count":       buy_count,
        "sell_count":      sell_count,
    }

    logger.info(f"Metrics: {total} alerts | {success_rate:.0f}% success | avg {avg_latency:.2f}s")
    return metrics


print("track_alert_metrics() defined")

def summarize_alerts(state: dict) -> dict:
    """
    Convert alert history to a pandas DataFrame and a printed summary.

    Parameters
    ----------
    state : dict
        Current alert system state.

    Returns
    -------
    dict
        Contains 'dataframe' (pd.DataFrame) and 'summary_text' (str).

    Examples
    --------
    >>> result = summarize_alerts(state)
    >>> len(result['dataframe'])
    20
    """
    history_list = list(state["history"])

    if not history_list:
        logger.warning("No alerts in history yet.")
        return {"dataframe": pd.DataFrame(), "summary_text": "No alerts yet."}

    df = pd.DataFrame(history_list)

    # Keep only readable columns
    display_cols = ["timestamp", "symbol", "action", "price", "quantity", "value", "status"]
    df = df[[c for c in display_cols if c in df.columns]]

    # Format numbers nicely
    if "price" in df.columns:
        df["price"] = df["price"].apply(lambda x: f"${x:.2f}")
    if "value" in df.columns:
        df["value"] = df["value"].apply(lambda x: f"${x:,.2f}")

    metrics = track_alert_metrics(state)
    summary_text = (
        f"Total Alerts : {metrics['total_alerts']}\n"
        f"Sent         : {metrics['sent_count']}  |  Failed: {metrics['failed_count']}\n"
        f"Success Rate : {metrics['success_rate']}%\n"
        f"BUY / SELL   : {metrics['buy_count']} / {metrics['sell_count']}\n"
        f"Total Value  : ${metrics['total_value_usd']:,.2f}"
    )

    logger.info(f"Summary generated for {len(df)} alerts.")
    return {"dataframe": df, "summary_text": summary_text}


print("summarize_alerts() defined")

def simulate_trade_signals(n: int = 20) -> List[dict]:
    """
    Generate n realistic fake trade signals for testing.

    Uses real stock tickers with approximate price ranges.
    Randomly assigns BUY/SELL and quantities.

    Parameters
    ----------
    n : int, optional
        Number of trade signals to generate, defaults to 20.

    Returns
    -------
    List[dict]
        List of dicts with symbol, action, price, quantity.

    Examples
    --------
    >>> signals = simulate_trade_signals(5)
    >>> signals[0].keys()
    dict_keys(['symbol', 'action', 'price', 'quantity'])
    """
    np.random.seed(42)  # Seed for reproducible results

    # Approximate real prices for well-known stocks (2025)
    stocks = {
        "AAPL": 189, "MSFT": 415, "GOOGL": 175,
        "AMZN": 195, "TSLA": 245, "NVDA": 875,
        "META": 520, "NFLX": 680,
    }

    signals = []
    for _ in range(n):
        symbol   = random.choice(list(stocks.keys()))
        base     = stocks[symbol]
        price    = round(base + np.random.normal(0, base * 0.02), 2)  # +-2% noise
        action   = random.choice(["BUY", "SELL"])
        quantity = random.randint(1, 100)
        signals.append({"symbol": symbol, "action": action,
                         "price": price, "quantity": quantity})

    logger.info(f"Generated {n} simulated trade signals.")
    return signals


# Quick test — display as table
test_signals = simulate_trade_signals(5)
pd.DataFrame(test_signals)

# -- DEMO CONFIG ---------------------------------------------------------------
DRY_RUN  = False   # Change to False to actually send real emails
N_ALERTS = 2     # How many trade signals to process
# ------------------------------------------------------------------------------

# 1. Initialize state
state   = create_alert_state(RECIPIENT_EMAIL)

# 2. Generate signals
signals = simulate_trade_signals(N_ALERTS)

print(f"{'='*55}")
print(f"  Running demo: {N_ALERTS} trade signals | DRY_RUN={DRY_RUN}")
print(f"{'='*55}\n")

# 3. Format + (optionally) send each alert
for i, sig in enumerate(signals, 1):
    alert = format_trade_alert(
        state,
        symbol   = sig["symbol"],
        action   = sig["action"],
        price    = sig["price"],
        quantity = sig["quantity"],
    )

    if DRY_RUN:
        # Simulate a successful send without actually emailing
        alert["status"]       = "sent"
        state["sent_count"]  += 1
        state["latencies"].append(round(random.uniform(0.3, 1.5), 2))
        state["history"].append(alert.copy())
        print(f"[{i:02d}] DRY RUN | {alert['subject']}")
    else:
        state = send_email_alert(state, alert, SMTP_CONFIG)
        print(f"[{i:02d}] {alert['status'].upper():6s} | {alert['subject']}")

print(f"\nDone! {state['sent_count']} sent, {state['failed_count']} failed.")

# Display the alert history as a formatted table
result = summarize_alerts(state)

print("Alert Summary")
print("-" * 40)
print(result["summary_text"])
print("\nFull History Table:")
display(result["dataframe"])

# Build a numeric dataframe for plotting (before price formatting)
history_list = list(state["history"])
df_raw = pd.DataFrame(history_list)

fig, axes = plt.subplots(1, 3, figsize=(16, 5))
fig.suptitle("Trade Alert System — Demo Dashboard", fontsize=14, fontweight="bold")

# Chart 1: Alert count per symbol
ax1 = axes[0]
symbol_counts = df_raw["symbol"].value_counts()
dominant_action = df_raw.groupby("symbol")["action"].agg(lambda x: x.mode()[0])
bar_colors = ["#2196F3" if dominant_action[s] == "BUY" else "#F44336"
              for s in symbol_counts.index]
symbol_counts.plot(kind="bar", ax=ax1, color=bar_colors, edgecolor="white")
ax1.set_title("Alert Volume by Stock", fontweight="bold")
ax1.set_xlabel("Symbol")
ax1.set_ylabel("Number of Alerts")
ax1.tick_params(axis="x", rotation=0)
ax1.grid(axis="y", alpha=0.3)
ax1.legend(handles=[Patch(color="#2196F3", label="Mostly BUY"),
                    Patch(color="#F44336", label="Mostly SELL")], fontsize=8)

# Chart 2: Price distribution BUY vs SELL
ax2 = axes[1]
buy_prices  = df_raw[df_raw["action"] == "BUY"]["price"]
sell_prices = df_raw[df_raw["action"] == "SELL"]["price"]
ax2.boxplot([buy_prices, sell_prices], labels=["BUY", "SELL"],
            patch_artist=True,
            boxprops=dict(facecolor="#E3F2FD"),
            medianprops=dict(color="#1565C0", linewidth=2))
ax2.set_title("Price Distribution: BUY vs SELL", fontweight="bold")
ax2.set_xlabel("Action")
ax2.set_ylabel("Price (USD)")
ax2.grid(axis="y", alpha=0.3)

# Chart 3: Success vs Failed pie
ax3 = axes[2]
metrics = track_alert_metrics(state)
sizes   = [metrics["sent_count"], metrics["failed_count"]]
labels  = [f"Sent ({metrics['sent_count']})", f"Failed ({metrics['failed_count']})"]
ax3.pie(sizes, labels=labels, autopct="%1.0f%%", startangle=90,
        colors=["#4CAF50", "#F44336"], explode=(0.05, 0.05),
        wedgeprops=dict(edgecolor="white", linewidth=2))
ax3.set_title("Email Delivery Rate", fontweight="bold")

plt.tight_layout()
plt.show()

# Bonus: Trade value over time (line plot)
fig, ax = plt.subplots(figsize=(14, 4))

buy_df  = df_raw[df_raw["action"] == "BUY"].reset_index(drop=True)
sell_df = df_raw[df_raw["action"] == "SELL"].reset_index(drop=True)

ax.plot(buy_df.index,  buy_df["value"],  marker="o", color="#2196F3",
        label="BUY alerts",  linewidth=1.5, markersize=5)
ax.plot(sell_df.index, sell_df["value"], marker="s", color="#F44336",
        label="SELL alerts", linewidth=1.5, markersize=5)

ax.set_title("Alert Trade Value Over Time (BUY vs SELL)", fontweight="bold")
ax.set_xlabel("Alert Index")
ax.set_ylabel("Trade Value (USD)")
ax.legend()
ax.grid(alpha=0.3)
ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f"${x:,.0f}"))
plt.tight_layout()
plt.show()

print("Testing edge cases...\n")

edge_cases = [
    {"symbol": "AAPL",  "action": "BUY",  "price": 0.01,   "quantity": 1,    "note": "Very low price"},
    {"symbol": "TSLA",  "action": "SELL", "price": 999.99, "quantity": 1000, "note": "Large block order"},
    {"symbol": "GOOGL", "action": "buy",  "price": 175.0,  "quantity": 5,    "note": "Lowercase 'buy' input"},
    {"symbol": "MSFT",  "action": "SELL", "price": 415.0,  "quantity": 0,    "note": "Zero quantity"},
]

edge_state = create_alert_state("edge_test@example.com")

for case in edge_cases:
    alert = format_trade_alert(edge_state, case["symbol"], case["action"],
                               case["price"], case["quantity"])
    alert["status"] = "sent"
    edge_state["sent_count"] += 1
    edge_state["history"].append(alert.copy())
    print(f"  Case   : {case['note']}")
    print(f"  Subject: {alert['subject']}")
    print(f"  Value  : ${alert['value']:,.2f}\n")

print("All edge cases handled without errors!")