#!/usr/bin/env python3
"""
Retrain emotion LoRA on complete dataset.
"""

import json
import random
import os
import shutil
import torch
from pathlib import Path
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, EarlyStoppingCallback
from peft import LoraConfig, get_peft_model, TaskType
from datasets import Dataset
from sklearn.metrics import f1_score, accuracy_score

# Load emotion data from labeled samples
emotion_data = []

# Load from labeled files that have emotion info
for fname in ["labeled_verified.jsonl", "labeled_expanded.jsonl", "labeled_real_world.jsonl"]:
    path = Path(f"/mnt/dolphinng5_predict/sentiment_engine/data/{fname}")
    if path.exists():
        with open(path) as f:
            for line in f:
                try:
                    item = json.loads(line.strip())
                    labels = item.get("labels", {})
                    text = item.get("text", item.get("raw_text", ""))
                    if text and labels.get("emotions"):
                        # Convert emotions to multi-label
                        emotion_labels = ["joy", "fear", "anger", "greed", "sadness", "neutral"]
                        multi_label = [0.0] * 6
                        for emo in labels["emotions"]:
                            if emo in emotion_labels:
                                multi_label[emotion_labels.index(emo)] = 1.0
                        if sum(multi_label) > 0:
                            emotion_data.append({"text": text, "labels": multi_label})
                except Exception as e:
                    pass

# Also add synthetic emotion data from templates
EMOTION_TEMPLATES = {
    "greed": ["FOMO driving {asset} to ${price}", "Buy the dip on {asset}! Loading bags", "Whale buying {asset} aggressively", "All in on {asset}! Diamond hands"],
    "fear": ["Panic selling {asset} at ${price}", "Major hack drains {asset} liquidity", "SEC crackdown sends {asset} plummeting", "Support broken on {asset}"],
    "joy": ["{asset} hits new ATH at ${price}! To the moon!", "ETF approved for {asset}!", "Massive gains on {asset}!"],
    "anger": ["Rug pull on {asset}! Devs drained liquidity!", "Exchange froze {asset} withdrawals again!", "Market manipulation on {asset}!"],
    "sadness": ["Lost life savings on {asset} crash", "Bag holder on {asset}... down 90%", "Rekt on {asset} leverage"],
    "neutral": ["{asset} consolidates at ${price}", "Low volume on {asset} at ${price}", "Market choppy for {asset}"],
}

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"]

# Generate synthetic emotion data
emotion_labels = ["joy", "fear", "anger", "greed", "sadness", "neutral"]

for emotion, templates in EMOTION_TEMPLATES.items():
    for template in templates:
        for _ in range(5):
            asset = random.choice(ASSETS)
            price = random.randint(100, 100000)
            text = template.format(asset=asset, price=price)
            labels = [0.0] * 6
            labels[emotion_labels.index(emotion)] = 1.0
            emotion_data.append({"text": text, "labels": labels})

# Also add from labeled data
for fname in ["labeled_verified.jsonl", "labeled_expanded.jsonl", "labeled_real_world.jsonl"]:
    path = Path(f"/mnt/dolphinng5_predict/sentiment_engine/data/{fname}")
    if path.exists():
        with open(path) as f:
            for line in f:
                try:
                    item = json.loads(line.strip())
                    labels = item.get("labels", {})
                    text = item.get("text", item.get("raw_text", ""))
                    if text and labels.get("emotions"):
                        multi_label = [0.0] * 6
                        for emo in labels["emotions"]:
                            if emo in emotion_labels:
                                multi_label[emotion_labels.index(emo)] = 1.0
                        if sum(multi_label) > 0:
                            emotion_data.append({"text": text, "labels": multi_label})
                except Exception as e:
                    pass

print(f"Total emotion samples: {len(emotion_data)}")

# Split
random.shuffle(emotion_data)
split = int(0.9 * len(emotion_data))
train_data = emotion_data[:split]
val_data = emotion_data[split:]

print(f"Train: {len(train_data)} | Val: {len(val_data)}")

# Label mapping
label2id = {l: i for i, l in enumerate(emotion_labels)}
id2label = {i: l for i, l in enumerate(emotion_labels)}

# Dataset
train_ds = Dataset.from_dict({"text": [d["text"] for d in train_data], "labels": [d["labels"] for d in train_data]})
val_ds = Dataset.from_dict({"text": [d["text"] for d in val_data], "labels": [d["labels"] for d in val_data]})

tokenizer = AutoTokenizer.from_pretrained("j-hartmann/emotion-english-distilroberta-base")

def tokenize(batch):
    return tokenizer(batch["text"], truncation=True, max_length=128, padding="max_length")

train_ds = train_ds.map(lambda b: tokenizer(b["text"], truncation=True, max_length=128, padding="max_length"), batched=True)
val_ds = val_ds.map(lambda b: tokenizer(b["text"], truncation=True, max_length=128, padding="max_length"), batched=True)
train_ds.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
val_ds.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

# Model + LoRA
model = AutoModelForSequenceClassification.from_pretrained(
    "j-hartmann/emotion-english-distilroberta-base",
    ignore_mismatched_sizes=True,
    num_labels=6,
    id2label=id2label,
    label2id=label2id,
    problem_type="multi_label_classification",
)

lora_config = LoraConfig(
    r=8, lora_alpha=16, lora_dropout=0.1,
    target_modules=("query", "key", "value", "intermediate.dense", "output.dense"),
    bias="none", task_type=TaskType.SEQ_CLS
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# Weighted loss
emotion_labels_list = ["joy", "fear", "anger", "greed", "sadness", "neutral"]
crypto_weights = {l: 1.0 for l in emotion_labels_list}
crypto_weights["greed"] = 2.0
crypto_weights["fear"] = 2.0
crypto_weights["joy"] = 1.5
class_weights = torch.tensor([crypto_weights[l] for l in emotion_labels_list], dtype=torch.float)

class WeightedTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        loss_fct = torch.nn.BCEWithLogitsLoss(pos_weight=class_weights.to(logits.device))
        loss = loss_fct(logits, labels.float())
        return (loss, outputs) if return_outputs else loss

output_dir = "./models/lora-distilroberta-crypto-emotion-v2"
os.makedirs(output_dir, exist_ok=True)
if os.path.exists(output_dir):
    shutil.rmtree(output_dir)

training_args = TrainingArguments(
    output_dir=output_dir,
    num_train_epochs=5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=16,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    warmup_ratio=0.1,
    weight_decay=0.01,
    max_grad_norm=1.0,
    eval_strategy="steps",
    eval_steps=25,
    save_strategy="steps",
    save_steps=25,
    save_total_limit=1,
    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,
    remove_unused_columns=False,
    report_to="none",
    seed=42,
)

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = (torch.sigmoid(torch.tensor(logits)) > 0.5).int().numpy()
    return {
        "f1_macro": f1_score(labels, preds, average="macro", zero_division=0),
        "f1_micro": f1_score(labels, preds, average="micro", zero_division=0),
        "accuracy": (preds == labels).mean(),
    }

trainer = WeightedTrainer(
    model=model,
    args=training_args,
    train_dataset=Dataset.from_dict({
        "input_ids": tokenizer([d["text"] for d in train_data], truncation=True, max_length=128, padding="max_length")["input_ids"],
        "attention_mask": tokenizer([d["text"] for d in train_data], truncation=True, max_length=128, padding="max_length")["attention_mask"],
        "labels": [d["labels"] for d in train_data],
    }),
    eval_dataset=Dataset.from_dict({
        "input_ids": tokenizer([d["text"] for d in val_data], truncation=True, max_length=128, padding="max_length")["input_ids"],
        "attention_mask": tokenizer([d["text"] for d in val_data], truncation=True, max_length=128, padding="max_length")["attention_mask"],
        "labels": [d["labels"] for d in val_data],
    }),
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3, early_stopping_threshold=0.001)],
)

print("🏋️ Training emotion LoRA v2...")
trainer.train()

best_path = "./models/lora-distilroberta-crypto-emotion-v2/best"
trainer.save_model(best_path)
AutoTokenizer.from_pretrained("j-hartmann/emotion-english-distilroberta-base").save_pretrained(best_path)
print(f"✅ Saved to {best_path}")

# Test
print("\n🧪 Testing emotion model...")
model.eval()
labels = ["joy", "fear", "anger", "greed", "sadness", "neutral"]

test_texts = [
    "BTC breaks 100k! New ATH, to the moon!",
    "Major hack on DeFi protocol, 50M drained",
    "Panic selling BTC at 50k, liquidation cascade",
    "FOMO buying ETH at 3k, loading bags",
    "Rug pull suspected, dev wallet drained liquidity",
    "BTC consolidates at 50k, no clear direction",
    "SEC sues exchange, regulatory crackdown",
    "ETF approved! Celebration time!",
]

from peft import PeftModel
model = PeftModel.from_pretrained(
    AutoModelForSequenceClassification.from_pretrained("j-hartmann/emotion-english-distilroberta-base", ignore_mismatched_sizes=True, num_labels=6, problem_type="multi_label_classification"),
    "./models/lora-distilroberta-crypto-emotion-v2/best"
)
model.eval()

for text in test_texts:
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.sigmoid(logits)[0]
        active = [(labels[i], probs[i].item()) for i in range(6) if probs[i] > 0.3]
        print(f'{text[:50]}')
        print(f'  Active: {active}')
        print()

label_map = {0: "Bearish", 1: "Bullish", 2: "Neutral"}