#!/usr/bin/env python3
"""
Augment the labeled set with carefully crafted crypto-specific samples
to balance classes and improve model performance.
"""

import json
import random
from pathlib import Path

# Load base set
with open("/mnt/dolphinng5_predict/sentiment_engine/data/final_labeled_set.jsonl") as f:
    base = [json.loads(line) for line in f]

print(f"Base samples: {len(base)}")

# Carefully crafted augmentation templates for each sentiment
TEMPLATES = {
    "Bullish": [
        # Real crypto bullish patterns
        "{asset} breaks resistance at ${price} with massive volume, institutional buyers stepping in",
        "{asset} surges to new ATH at ${price} as {catalyst} drives inflows",
        "Institutional adoption drives {asset} to ${price}, whale accumulation evident",
        "ETF approval sends {asset} to ${price}, massive inflows expected",
        "{asset} breaks out of consolidation at ${price}, next target ${target}",
        "Major partnership announced for {asset}, price surges to ${price}",
        "Whale accumulation pushes {asset} above ${price}, on-chain metrics bullish",
        "DeFi protocol {asset} TVL hits record high at ${price}",
        "Layer 2 adoption drives {asset} to ${price}, scaling solution working",
        "Staking rewards increase for {asset}, yield hunters accumulate at ${price}",
        "Major exchange lists {asset}, price jumps to ${price}",
        "Regulatory clarity for {asset} drives price to ${price}",
        "Upgrade activates for {asset}, scaling improves, price to ${price}",
        "Cross-chain bridge launches for {asset}, liquidity flows at ${price}",
        "HODL strong hands, diamond hands win as {asset} holds ${price}",
    ],
    "Bearish": [
        # Real crypto bearish patterns
        "Major hack on {asset} protocol drains ${amount}M, price crashes to ${price}",
        "SEC sues {asset} team for unregistered securities, price drops to ${price}",
        "Rug pull suspected on {asset}, dev wallet drains liquidity, price to ${price}",
        "Exchange delists {asset}, panic selling drives price to ${price}",
        "Regulatory crackdown on {asset} sends price plummeting to ${price}",
        "Massive liquidation cascade wipes {asset} longs, price drops to ${price}",
        "Support broken on {asset} at ${price}, bearish continuation expected",
        "Whale dumping {asset}, massive sell wall at ${price}",
        "Ransomware attackers dump {asset} for BTC, price crashes to ${price}",
        "Liquidity pulled from {asset} pools, price collapses to ${price}",
        "51% attack feared on {asset} as hashrate drops, price to ${price}",
        "Smart contract exploit on {asset}, ${amount}M stolen, price to ${price}",
        "Market manipulation suspected on {asset}, coordinated dump to ${price}",
        "Exchange halts {asset} withdrawals, panic selling to ${price}",
        "Stablecoin depeg triggers {asset} selloff to ${price}",
    ],
    "Neutral": [
        # Real neutral/consolidation patterns
        "{asset} consolidates at ${price} in tight range, awaiting catalyst",
        "Low volume on {asset} at ${price}, market awaiting direction",
        "Sideways action on {asset} at ${price}, no clear direction",
        "{asset} forms doji at ${price}, direction unclear",
        "Range-bound trading for {asset} between ${low} and ${high}",
        "Accumulation phase for {asset} around ${price}",
        "Low volatility on {asset} at ${price}, volume drying up",
        "Market in wait-and-see mode for {asset} at ${price}",
        "{asset} at ${price} with mixed on-chain signals",
        "No fresh news on {asset}, price stable at ${price}",
        "Choppy action for {asset} at ${price}, traders cautious",
        "{asset} forms pennant at ${price}, breakout direction unknown",
    ],
}

ASSETS = ["BTC", "ETH", "SOL", "AVAX", "MATIC", "DOT", "LINK", "ARB", "OP", "NEAR", "FET", "STX", "ZIL", "XTZ", "ENJ", "ETC", "TRX", "LTC", "DASH", "ONG", "ONE", "ALGO", "DOGE", "XLM", "ATOM", "KSM", "APT", "SUI", "ICP", "QNT", "INJ"]

CATALYSTS = [
    "institutional inflows", "ETF approval", "whale accumulation", "DeFi adoption",
    "institutional custody", "staking rewards", "protocol upgrade", "cross-chain bridge",
    "major partnership", "exchange listing", "regulatory clarity", "TVL growth"
]

def augment():
    with open("/mnt/dolphinng5_predict/sentiment_engine/data/final_labeled_set.jsonl") as f:
        base = [json.loads(line) for line in open("/mnt/dolphinng5_predict/sentiment_engine/data/final_labeled_set.jsonl")]
    
    augmented = list(base)  # Start with base
    
    for sentiment, templates in TEMPLATES.items():
        # Generate more samples for underrepresented classes
        target = 200 if sentiment == "Bearish" else 150 if sentiment == "Bullish" else 100
        current = len([d for d in base if d["sentiment"] == sentiment])
        needed = max(0, target - current)
        
        if needed > 0:
            print(f"Generating {needed} {sentiment} samples...")
            for _ in range(needed):
                template = random.choice(templates)
                asset = random.choice(ASSETS)
                price = random.randint(100, 100000)
                target_price = price + random.randint(100, 5000)
                low = price - random.randint(10, 500)
                high = price + random.randint(10, 500)
                amount = random.randint(5, 200)
                catalyst = random.choice(CATALYSTS)
                
                text = template.format(
                    asset=asset, price=price, target=target_price,
                    low=low, high=high, amount=amount, catalyst=catalyst
                )
                
                augmented.append({
                    "text": text,
                    "sentiment": sentiment,
                    "event_type": "price_action",
                    "entities": [asset],
                    "source": f"augmented_{sentiment.lower()}",
                })
    
    # Shuffle and save
    random.shuffle(augmented)
    print(f"Total augmented samples: {len(augmented)}")
    
    # Count by sentiment
    from collections import Counter
    dist = Counter(d["sentiment"] for d in augmented)
    print(f"Distribution: {dict(dist)}")
    
    with open("/mnt/dolphinng5_predict/sentiment_engine/data/final_augmented_set.jsonl", "w") as f:
        for d in augmented:
            json.dump(d, f)
            f.write("\n")
    
    print("Saved to final_augmented_set.jsonl")

if __name__ == "__main__":
    import json
    augment()