#!/usr/bin/env python3
"""
Ultra-fast FinBERT fine-tuning demo (CPU, ~10 min).
Uses tiny dataset, 1 epoch, aggressive settings for demo purposes.
"""

import json
import random
import torch
import torch.nn as nn
import numpy as np
from pathlib import Path
from typing import List, Dict
from torch.utils.data import Dataset
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    TrainingArguments, Trainer, EarlyStoppingCallback
)
from datasets import load_dataset
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score
from sklearn.utils.class_weight import compute_class_weight
import torch.nn as np

SENTIMENT_LABELS = ["Bearish", "Bullish", "Neutral"]
SENTIMENT_MAP = {"Bearish": 0, "Bullish": 1, "Neutral": 2}

# Minimal real samples
REAL_SAMPLES = [
    {"text": "BTC surges to new all-time high as institutional adoption accelerates!", "label_id": 1},
    {"text": "ETH breaks $4000 resistance with massive volume!", "label_id": 1},
    {"text": "Institutional adoption drives Bitcoin higher!", "label_id": 1},
    {"text": "Bitcoin breaks $100k! New ATH!", "label_id": 1},
    {"text": "Institutional inflows hit record high", "label_id": 1},
    ("BTC crashes 50% in hours", 0),
    ("Exchange hacked, $100M stolen", 0),
    ("SEC sues major exchange", 0),
    ("Bitcoin crashes hard, panic selling everywhere", 0),
    ("Massive liquidation cascade wipes out $200M in longs", 0),
    ("Whale sells 10000 BTC", 0),
    ("Bitcoin price drops 50%", 0),
    ("Support broken with bearish structure", 0),
    ("Panic selling and forced liquidation", 0),
    ("BTC at $50k, ETH at $3k", 2),
    ("Market consolidating in range", 2),
    ("Bitcoin remains stable around $30k", 2),
    ("Market consolidating with no clear direction", 2),
]

SENTIMENT_LABELS = ["Bearish", "Bullish", "Neutral"]
SENTIMENT_MAP = {"Bearish": 0, "Bullish": 1, "Neutral": 2}

# Real crypto events
REAL_EVENTS = [
    {"text": "XRP bridge drained for $200,000 after software mistook fake deposits for real ones.", "label_id": 0},
    {"text": "Major hack on DeFi protocol drains $50M. Users panic as TVL collapses.", "label_id": 0},
    {"text": "KuCoin Lists Catizen (CATI) for Spot Trading.", "label_id": 1},
    {"text": "Binance Becomes First Exchange to List Trump-Linked WLFI Token.", "label_id": 1},
    {"text": "SEC files lawsuit against major exchange for unregistered securities.", "label_id": 0},
    {"text": "CFTC files to dismiss CME's lawsuit over crypto perpetual futures.", "label_id": 2},
    {"text": "Ethereum Dencun upgrade activates Proto-Danksharding (EIP-4844).", "label_id": 1},
    {"text": "Ethereum Shanghai upgrade goes live. Stakers can now withdraw.", "label_id": 1},
    {"text": "JPMorganChase and Coinbase Launch Strategic Partnership.", "label_id": 1},
    {"text": "Chainlink and Mastercard Partner to Enable Over 3 Billion Cardholders.", "label_id": 1},
    {"text": "Bitcoin whale moves $116 million in BTC after 11-year dormancy.", "label_id": 2},
    {"text": "Breaking: Fed pauses rate hikes. Bitcoin jumps 5% on dovish pivot.", "label_id": 1},
    {"text": "Massive liquidation cascade wipes out $200M in longs.", "label_id": 0},
    {"text": "Bitcoin ETF inflows hit $731M, highest since January.", "label_id": 1},
    {"text": "Coinbase delists XRP after SEC lawsuit.", "label_id": 0},
]

SENTIMENT_LABELS = ["Bearish", "Bullish", "Neutral"]
SENTIMENT_MAP = {"Bearish": 0, "Bullish": 1, "Neutral": 2}

class QuickDataset(torch.utils.data.Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=64):
        self.texts = texts
        self.labels = labels
        self.tokenizer = AutoTokenizer.from_pretrained("ProsusAI/finbert")
        self.max_len = 64
    
    def __len__(self): return len(self.texts)
    def __getitem__(self, i):
        enc = self.tokenizer(self.texts[i], truncation=True, max_length=self.max_len, 
                            padding="max_length", return_tensors="pt")
        return {
            "input_ids": enc["input_ids"].squeeze(0),
            "attention_mask": enc["attention_mask"].squeeze(0),
            "labels": torch.tensor(self.labels[i], dtype=torch.long)
        }

def main():
    print("=" * 50)
    print("Quick FinBERT Crypto Fine-Tune (CPU, ~5 min)")
    print("=" * 50)
    
    # Build tiny dataset
    texts = []
    labels = []
    
    # Manual samples
    for text, label in [
        ("BTC breaks $100k! New ATH!", 1),
        ("ETH to $10k by EOY, accumulate now", 1),
        ("Institutional inflows hit record high", 1),
        ("Bitcoin reaches new all-time high", 1),
        ("Ethereum merge successful, staking rewards now live", 1),
        ("Massive ETF inflows drive Bitcoin to new highs", 1),
        ("Golden cross confirmed on Bitcoin weekly chart", 1),
        ("Institutional adoption drives Bitcoin higher", 1),
        ("ETF approval drives massive inflows", 1),
        ("Market is bullish on Bitcoin", 1),
        ("BTC crashes 50% in hours", 0),
        ("Exchange hacked, $100M stolen", 0),
        ("SEC sues major exchange", 0),
        ("Bitcoin crashes hard, panic selling everywhere", 0),
        ("Massive liquidation cascade wipes out $200M in longs", 0),
        ("VIX drops below 15 as market volatility decreases", 0),
        ("Whale sells 10000 BTC", 0),
        ("Bitcoin price drops 50%", 0),
        ("Support broken with bearish structure", 0),
        ("Panic selling and forced liquidation", 0),
        ("BTC at $50k, ETH at $3k", 2),
        ("Market consolidating in range", 2),
        ("Bitcoin remains stable around $30k", 2),
        ("VIX drops below 15 as market volatility decreases", 2),
        ("Market consolidating with no clear direction", 2),
        ("Bitcoin price stable around $30k", 2),
        ("Consolidation phase continues", 2),
        ("Market in wait-and-see mode", 2),
        ("Sideways action continues", 2),
        ("Low volatility environment persists", 2),
    ]:
        texts.append(t)
        labels.append(l)
    
    # Add real events
    for event in [
        {"text": "XRP bridge drained for $200,000 after software mistook fake deposits.", "label_id": 0},
        {"text": "Major hack on DeFi protocol drains $50M.", "label_id": 0},
        {"text": "KuCoin Lists Catizen (CATI) for Spot Trading.", "label_id": 1},
        {"text": "Binance Becomes First Exchange to List Trump-Linked WLFI Token.", "label_id": 1},
        {"text": "SEC files lawsuit against major exchange for unregistered securities.", "label_id": 0},
        {"text": "Ethereum Dencun upgrade activates Proto-Danksharding (EIP-4844).", "label_id": 1},
        {"text": "JPMorganChase and Coinbase Launch Strategic Partnership.", "label_id": 1},
        {"text": "Chainlink and Mastercard Partner to Enable Over 3 Billion Cardholders.", "label_id": 1},
        {"text": "Bitcoin whale moves $116 million in BTC after 11-year dormancy.", "label_id": 2},
        {"text": "Breaking: Fed pauses rate hikes. Bitcoin jumps 5% on dovish pivot.", "label_id": 1},
        {"text": "Massive liquidation cascade wipes out $200M in longs.", "label_id": 0},
        {"text": "Bitcoin ETF inflows hit $731M, highest since January.", "label_id": 1},
        {"text": "Coinbase delists XRP after SEC lawsuit.", "label_id": 0},
    ]:
        texts.append(item["text"])
        labels.append(item["label_id"])
    
    # Augmented
    assets = ["BTC", "ETH", "SOL", "AVAX", "MATIC"]
    templates = {
        1: ["{a} surges to new highs", "{a} breaks resistance at ${p}", "Institutional adoption drives {a} higher"],
        0: ["{a} crashes {p}%", "{a} breaks support at ${p}", "Panic selling in {a}"],
        2: ["{a} consolidates at ${p}", "{a} trades sideways", "Market waits for {a} direction"],
    }
    for _ in range(500):
        sid = random.randint(0, 2)
        a = random.choice(["BTC", "ETH", "SOL", "AVAX", "MATIC"])
        t = random.choice([p for p in range(3) if p in [0,1,2]])  # simplified
        template = random.choice(templates[sid])
        text = template.format(a=random.choice(assets), p=random.randint(100,100000))
        texts.append(text)
        labels.append(sid)
    
    print(f"Total samples: {len(texts)}")
    
    # Split
    from sklearn.model_selection import train_test_split
    train_t, temp_t, train_l, temp_l = train_test_split(texts, labels, test_size=0.3, random_state=42, stratify=labels)
    temp_t, test_t, temp_l, test_l = train_test_split(temp_t, temp_l, test_size=0.5, random_state=42, stratify=temp_l)
    
    print(f"Train: {len(train_t)}, Val: {len(temp_t)}, Test: {len(test_t)}")
    
    # Tokenizer & Model
    from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, EarlyStoppingCallback
    import torch
    
    tokenizer = AutoTokenizer.from_pretrained("ProsusAI/finbert")
    model = AutoModelForSequenceClassification.from_pretrained(
        "ProsusAI/finbert", num_labels=3,
        id2label={0:"Bearish",1:"Bullish",2:"Neutral"},
        label2id={"Bearish":0,"Bullish":1,"Neutral":2}
    )
    
    class QuickDataset(torch.utils.data.Dataset):
        def __init__(self, texts, labels, tokenizer, max_len=64):
            self.texts = texts; self.labels = labels
            self.tokenizer = AutoTokenizer.from_pretrained("ProsusAI/finbert")
            self.max_len = 64
        def __len__(self): return len(self.texts)
        def __getitem__(self, i):
            enc = self.tokenizer(self.texts[i], truncation=True, max_length=64, padding="max_length", return_tensors="pt")
            return {"input_ids": enc["input_ids"].squeeze(0), "attention_mask": enc["attention_mask"].squeeze(0), "labels": torch.tensor(self.labels[i], dtype=torch.long)}
    
    train_ds = QuickDataset(texts[:len(train_t)], labels[:len(train_t)], None)
    val_ds = QuickDataset(texts[len(train_t):len(train_t)+len(temp_t)], labels[len(train_l):len(train_l)+len(temp_l)], None)
    test_ds = QuickDataset(texts[-len(test_t):], labels[-len(test_l):], None)
    
    # Fix: create datasets properly
    train_texts = texts[:len(train_t)]
    train_labels = labels[:len(train_l)]
    val_texts = texts[len(train_t):len(train_t)+len(temp_t)]
    val_labels = labels[len(train_l):len(train_l)+len(temp_l)]
    test_texts = texts[-len(test_t):]
    test_labels = labels[-len(test_l):]
    
    train_ds = QuickDataset(train_texts, train_labels, None)
    val_ds = QuickDataset(val_texts, val_labels, None)
    test_ds = QuickDataset(test_texts, test_labels, None)
    
    from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, EarlyStoppingCallback
    import torch
    
    tokenizer = AutoTokenizer.from_pretrained("ProsusAI/finbert")
    model = AutoModelForSequenceClassification.from_pretrained(
        "ProsusAI/finbert", num_labels=3,
        id2label={0:"Bearish",1:"Bullish",2:"Neutral"},
        label2id={"Bearish":0,"Bullish":1,"Neutral":2}
    )
    
    class QuickDataset(torch.utils.data.Dataset):
        def __init__(self, texts, labels, tokenizer, max_len=64):
            self.texts = texts; self.labels = labels
            self.tokenizer = tokenizer; self.max_len = 64
        def __len__(self): return len(self.texts)
        def __getitem__(self, i):
            enc = self.tokenizer(self.texts[i], truncation=True, max_length=self.max_len, padding="max_length", return_tensors="pt")
            return {"input_ids": enc["input_ids"].squeeze(0), "attention_mask": enc["attention_mask"].squeeze(0), "labels": torch.tensor(self.labels[i], dtype=torch.long)}
    
    train_ds = QuickDataset(train_texts, train_labels, tokenizer)
    val_ds = QuickDataset(val_texts, val_labels, tokenizer)
    test_ds = QuickDataset(test_texts, test_labels, tokenizer)
    
    # Train
    trainer = Trainer(
        model=AutoModelForSequenceClassification.from_pretrained("ProsusAI/finbert", num_labels=3, id2label={0:"Bearish",1:"Bullish",2:"Neutral"}, label2id={"Bearish":0,"Bullish":1,"Neutral":2}),
        args=TrainingArguments(
            
            num_train_epochs=1,
            per_device_train_batch_size=16,
            per_device_eval_batch_size=32,
            gradient_accumulation_steps=2,
            warmup_ratio=0.1,
            learning_rate=2e-5,
            lr_scheduler_type="cosine",
            evaluation_strategy="epoch",
            save_strategy="epoch",
            load_best_model_at_end=True,
            metric_for_best_model="f1_macro",
            greater_is_better=True,
            fp16=False,
            dataloader_num_workers=0,
            logging_steps=10,
            save_total_limit=1,
            remove_unused_columns=False,
            report_to="none",
            
        ),
        train_dataset=QuickDataset(train_texts, train_labels, tokenizer, max_len=64),
        eval_dataset=QuickDataset(val_texts, val_labels, tokenizer, max_len=64),
        tokenizer=AutoTokenizer.from_pretrained("ProsusAI/finbert"),
        compute_metrics=lambda ep: {"f1_macro": f1_score(ep.label_ids, np.argmax(ep.predictions, axis=-1), average="macro")},
        callbacks=[EarlyStoppingCallback(early_stopping_patience=1)]
    )
    
    from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, EarlyStoppingCallback
    from sklearn.metrics import f1_score
    import torch
    
    trainer = Trainer(
        model=AutoModelForSequenceClassification.from_pretrained("ProsusAI/finbert", num_labels=3, id2label={0:"Bearish",1:"Bullish",2:"Neutral"}, label2id={"Bearish":0,"Bullish":1,"Neutral":2}),
        args=TrainingArguments(
            
            num_train_epochs=1,
            per_device_train_batch_size=16,
            per_device_eval_batch_size=32,
            gradient_accumulation_steps=2,
            warmup_ratio=0.1,
            learning_rate=2e-5,
            lr_scheduler_type="cosine",
            evaluation_strategy="epoch",
            save_strategy="epoch",
            load_best_model_at_end=True,
            metric_for_best_model="f1_macro",
            greater_is_better=True,
            fp16=False,
            dataloader_num_workers=0,
            logging_steps=10,
            save_total_limit=1,
            remove_unused_columns=False,
            report_to="none",
            
        ),
        train_dataset=QuickDataset(train_texts, train_labels, AutoTokenizer.from_pretrained("ProsusAI/finbert"), max_len=64),
        eval_dataset=QuickDataset(val_texts, val_labels, AutoTokenizer.from_pretrained("ProsusAI/finbert"), max_len=64),
        tokenizer=AutoTokenizer.from_pretrained("ProsusAI/finbert"),
        compute_metrics=lambda ep: {"f1_macro": f1_score(ep.label_ids, np.argmax(ep.predictions, axis=-1), average="macro")},
        callbacks=[EarlyStoppingCallback(early_stopping_patience=1)]
    )
    
    print("Training (1 epoch, ~2-3 min)...")
    trainer.train()
    
    # Test
    print("\nTest results:")
    results = trainer.evaluate(ep=ep) if False else trainer.evaluate()
    print(f"Test: {results}")
    
    trainer.save_model("./models/finbert-crypto-quick")
    AutoTokenizer.from_pretrained("ProsusAI/finbert").save_pretrained("./models/finbert-crypto-quick")
    print("Saved!")
    
    # Quick test
    model.eval()
    for text in ["BTC surges to new ATH!", "Bitcoin crashes 50%!", "BTC consolidates at $50k"]:
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=64, padding=True)
        with torch.no_grad():
            out = model(**inputs)
            probs = torch.softmax(out.logits, dim=-1)[0]
            pred = torch.argmax(probs).item()
            pol = probs[1].item() - probs[0].item()
            print(f"  '{text}' -> {['Bearish','Bullish','Neutral'][pred]} (pol: {pol:.3f})")
    print("Done!")

if __name__ == "__main__":
    import random, torch
    from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, EarlyStoppingCallback
    from sklearn.model_selection import train_test_split
    from sklearn.metrics import f1_score
    import numpy as np
    main()