# data/twitter_pipeline.py
import tweepy
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from models.finbert import FinBERTSentiment

class TwitterSentimentPipeline:
    def __init__(self, bearer_token: str):
        self.client = tweepy.Client(bearer_token=bearer_token)
        self.finbert = FinBERTSentiment()
        self.ticker_pattern = None

    def fetch_tweets(self, query: str, max_results: int = 100) -> pd.DataFrame:
        tweets = tweepy.Paginator(
            self.client.search_recent_tweets,
            query=query + " -is:retweet lang:en",
            tweet_fields=["created_at", "public_metrics", "author_id"],
            max_results=min(max_results, 100),
        ).flatten(limit=max_results)

        records = []
        for tweet in tweets:
            metrics = tweet.public_metrics
            records.append({
                "id": tweet.id,
                "text": tweet.text,
                "created_at": tweet.created_at,
                "likes": metrics["like_count"],
                "retweets": metrics["retweet_count"],
                "replies": metrics["reply_count"],
                "engagement": metrics["like_count"] + metrics["retweet_count"] * 2 + metrics["reply_count"],
            })
        return pd.DataFrame(records)

    def extract_tickers(self, text: str) -> list[str]:
        import re
        tickers = re.findall(r'\$([A-Z]{1,5})', text)
        return list(set(tickers))

    def compute_engagement_weight(self, row: dict) -> float:
        base = np.log1p(row["engagement"])
        return min(base / 10, 1.0)

    def process_batch(self, tweets: pd.DataFrame) -> dict:
        results = {}
        for _, row in tweets.iterrows():
            tickers = self.extract_tickers(row["text"])
            sentiment = self.finbert.analyze(row["text"])

            for ticker in tickers:
                if ticker not in results:
                    results[ticker] = []
                results[ticker].append({
                    "sentiment": sentiment["score"],
                    "confidence": sentiment["confidence"],
                    "engagement": row["engagement"],
                    "weight": self.compute_engagement_weight(row),
                    "timestamp": row["created_at"],
                })
        return results

    def aggregate_signals(self, ticker_data: list[dict]) -> dict:
        if not ticker_data:
            return {"score": 0, "volume": 0}

        scores = [d["sentiment"] for d in ticker_data]
        weights = [d["weight"] * d["confidence"] for d in ticker_data]

        weighted_score = np.average(scores, weights=weights)
        volume_z = len(ticker_data) / 100  # Normalize

        return {
            "score": weighted_score,
            "volume": volume_z,
            "n_posts": len(ticker_data),
            "avg_confidence": np.mean([d["confidence"] for d in ticker_data]),
        }

# backtest/engine.py
import pandas as pd
import numpy as np

class SentimentBacktester:
    def __init__(self, initial_capital: float = 100_000):
        self.capital = initial_capital
        self.transaction_cost = 0.001

    def backtest(self, signals: pd.DataFrame, returns: pd.DataFrame) -> dict:
        positions = signals.apply(lambda x: 1 if x > 0.5 else (-1 if x < -0.5 else 0))
        positions = positions.shift(1)  # Execute next period

        strategy_returns = (positions * returns).sum(axis=1)
        costs = positions.diff().abs().sum(axis=1) * self.transaction_cost
        net_returns = strategy_returns - costs

        cum_returns = (1 + net_returns).cumprod()
        total_return = cum_returns.iloc[-1] - 1
        annual_return = (1 + total_return) ** (252 / len(net_returns)) - 1
        annual_vol = net_returns.std() * np.sqrt(252)
        sharpe = annual_return / (annual_vol + 1e-10)

        peak = cum_returns.cummax()
        drawdown = (cum_returns - peak) / peak
        max_drawdown = drawdown.min()

        return {
            "total_return": total_return,
            "annual_return": annual_return,
            "annual_vol": annual_vol,
            "sharpe": sharpe,
            "max_drawdown": max_drawdown,
            "win_rate": (net_returns > 0).mean(),
            "avg_trades_per_day": positions.diff().abs().sum(axis=1).mean(),
        }

# main.py
from data.twitter_pipeline import TwitterSentimentPipeline
from backtest.engine import SentimentBacktester

pipeline = TwitterSentimentPipeline(bearer_token="YOUR_TOKEN")

tickers = ["AAPL", "MSFT", "TSLA", "NVDA", "AMZN"]
for ticker in tickers:
    tweets = pipeline.fetch_tweets(f"${ticker}", max_results=200)
    batch_results = pipeline.process_batch(tweets)
    signal = pipeline.aggregate_signals(batch_results.get(ticker, []))
    print(f"${ticker}: sentiment={signal['score']:.3f} volume={signal['volume']:.2f} posts={signal['n_posts']}")