#!/usr/bin/env python3
"""
Fine-tune existing models with newly labeled data from labeling pipeline.
Loads existing fine-tuned models and continues training on labeled_verified.jsonl
"""

import json
import torch
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 sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score

# ============================================================
# LABELS & CONSTANTS
# ============================================================

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

EMOTION_LABELS = ["joy", "fear", "anger", "greed", "sadness", "neutral"]
EMOTION_MAP = {l: i for i, l in enumerate(EMOTION_LABELS)}

EVENT_LABELS = [
    "listing", "delisting", "hack", "regulatory", "governance",
    "upgrade", "partnership", "earnings", "macro",
    "liquidation", "whale", "manipulation"
]
EVENT_MAP = {l: i for i, l in enumerate(EVENT_LABELS)}

# ============================================================
# LOAD LABELED DATA FROM LABELING PIPELINE
# ============================================================

def load_labeled_data(label_file: str):
    """Load verified labeled data from JSONL file"""
    sentiment_texts, sentiment_labels = [], []
    event_texts, event_labels = [], []
    emotion_texts, emotion_labels = [], []
    
    with open(label_file) as f:
        for line in f:
            r = json.loads(line)
            if r.get('verified', False):
                text = r['text']
                labels = r['labels']
                
                # Sentiment
                sentiment_texts.append(text)
                sentiment_labels.append(SENTIMENT_MAP[labels['sentiment']])
                
                # Event (single label)
                event_type = labels.get('event_type', 'listing')
                event_lbl = [0] * len(EVENT_LABELS)
                if event_type in EVENT_MAP:
                    event_lbl[EVENT_MAP[event_type]] = 1
                event_texts.append(text)
                event_labels.append(event_lbl)
                
                # Emotion (multi-label)
                emotions = labels.get('emotions', {})
                emotion_lbl = [0] * len(EMOTION_LABELS)
                for emotion, score in emotions.items():
                    if emotion in EMOTION_MAP and score > 0.5:
                        emotion_lbl[EMOTION_MAP[emotion]] = 1
                # If no emotions detected, set neutral
                if sum(emotion_lbl) == 0:
                    emotion_lbl[EMOTION_MAP['neutral']] = 1
                emotion_texts.append(text)
                emotion_labels.append(emotion_lbl)
    
    return {
        'sentiment': (sentiment_texts, sentiment_labels),
        'event': (event_texts, event_labels),
        'emotion': (emotion_texts, emotion_labels)
    }

# ============================================================
# DATASET CLASS
# ============================================================

class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=64, is_multilabel=False):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.is_multilabel = is_multilabel
    
    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"
        )
        lbl = self.labels[i]
        if self.is_multilabel:
            lbl = torch.tensor(lbl, dtype=torch.float)
        else:
            lbl = torch.tensor(lbl, dtype=torch.long)
        return {
            "input_ids": enc["input_ids"].squeeze(0),
            "attention_mask": enc["attention_mask"].squeeze(0),
            "labels": lbl
        }

# ============================================================
# TRAINING FUNCTIONS
# ============================================================

def compute_metrics_single(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return {"accuracy": accuracy_score(labels, predictions), "f1_macro": f1_score(labels, predictions, average="macro")}

def compute_metrics_multilabel(eval_pred):
    predictions, labels = eval_pred
    predictions = (np.array(predictions) > 0.5).astype(int)
    return {"f1_macro": f1_score(labels, predictions, average="macro")}

def fine_tune_sentiment(model_path: str, texts: List[str], labels: List[int]):
    print(f"\n{'='*50}")
    print("FINE-TUNING SENTIMENT (FinBERT)")
    print(f"{'='*50}")
    print(f"Training samples: {len(texts)}")
    
    train_t, val_t, train_l, val_l = train_test_split(
        texts, labels, test_size=0.2, random_state=42, stratify=labels
    )
    print(f"Train: {len(train_t)}, Val: {len(val_t)}")
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    
    train_ds = TextClassificationDataset(train_t, train_l, tokenizer, max_len=64)
    val_ds = TextClassificationDataset(val_t, val_l, tokenizer, max_len=64)
    
    trainer = Trainer(
        model=model,
        args=TrainingArguments(
            output_dir="./models/finbert-crypto-sentiment-ft",
            num_train_epochs=1,
            per_device_train_batch_size=8,
            per_device_eval_batch_size=16,
            gradient_accumulation_steps=4,
            warmup_ratio=0.1,
            learning_rate=1e-5,  # Lower LR for fine-tuning
            lr_scheduler_type="cosine",
            eval_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=5,
            save_total_limit=1,
            remove_unused_columns=False,
            report_to="none",
        ),
        train_dataset=train_ds,
        eval_dataset=val_ds,
        tokenizer=tokenizer,
        compute_metrics=compute_metrics_single,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=1)]
    )
    
    print("Training Sentiment (1 epoch)...")
    trainer.train()
    
    model.save_pretrained("./models/finbert-crypto-sentiment")
    tokenizer.save_pretrained("./models/finbert-crypto-sentiment")
    print("✅ Sentiment model fine-tuned and saved!")
    return model

def fine_tune_events(model_path: str, texts: List[str], labels: List[List[int]]):
    print(f"\n{'='*50}")
    print("FINE-TUNING EVENTS (BERT)")
    print(f"{'='*50}")
    print(f"Training samples: {len(texts)}")
    
    train_t, val_t, train_l, val_l = train_test_split(
        texts, labels, test_size=0.2, random_state=42
    )
    print(f"Train: {len(train_t)}, Val: {len(val_t)}")
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_path,
        num_labels=len(EVENT_LABELS),
        id2label={i:l for i,l in enumerate(EVENT_LABELS)},
        label2id=EVENT_MAP,
        problem_type="multi_label_classification",
        ignore_mismatched_sizes=True
    )
    
    train_ds = TextClassificationDataset(train_t, train_l, tokenizer, max_len=64, is_multilabel=True)
    val_ds = TextClassificationDataset(val_t, val_l, tokenizer, max_len=64, is_multilabel=True)
    
    trainer = Trainer(
        model=model,
        args=TrainingArguments(
            output_dir="./models/bert-crypto-events-ft",
            num_train_epochs=1,
            per_device_train_batch_size=8,
            per_device_eval_batch_size=16,
            gradient_accumulation_steps=4,
            warmup_ratio=0.1,
            learning_rate=1e-5,
            lr_scheduler_type="cosine",
            eval_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=5,
            save_total_limit=1,
            remove_unused_columns=False,
            report_to="none",
        ),
        train_dataset=train_ds,
        eval_dataset=val_ds,
        tokenizer=tokenizer,
        compute_metrics=compute_metrics_multilabel,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=1)]
    )
    
    print("Training Events (1 epoch)...")
    trainer.train()
    
    model.save_pretrained("./models/bert-crypto-events")
    tokenizer.save_pretrained("./models/bert-crypto-events")
    print("✅ Event model fine-tuned and saved!")
    return model

def fine_tune_emotion(model_path: str, texts: List[str], labels: List[List[int]]):
    print(f"\n{'='*50}")
    print("FINE-TUNING EMOTION (DistilRoBERTa)")
    print(f"{'='*50}")
    print(f"Training samples: {len(texts)}")
    
    train_t, val_t, train_l, val_l = train_test_split(
        texts, labels, test_size=0.2, random_state=42
    )
    print(f"Train: {len(train_t)}, Val: {len(val_t)}")
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_path,
        num_labels=len(EMOTION_LABELS),
        id2label={i:l for i,l in enumerate(EMOTION_LABELS)},
        label2id=EMOTION_MAP,
        problem_type="multi_label_classification",
        ignore_mismatched_sizes=True
    )
    
    train_ds = TextClassificationDataset(train_t, train_l, tokenizer, max_len=64, is_multilabel=True)
    val_ds = TextClassificationDataset(val_t, val_l, tokenizer, max_len=64, is_multilabel=True)
    
    trainer = Trainer(
        model=model,
        args=TrainingArguments(
            output_dir="./models/distilroberta-crypto-emotion-ft",
            num_train_epochs=1,
            per_device_train_batch_size=8,
            per_device_eval_batch_size=16,
            gradient_accumulation_steps=4,
            warmup_ratio=0.1,
            learning_rate=1e-5,
            lr_scheduler_type="cosine",
            eval_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=5,
            save_total_limit=1,
            remove_unused_columns=False,
            report_to="none",
        ),
        train_dataset=train_ds,
        eval_dataset=val_ds,
        tokenizer=tokenizer,
        compute_metrics=compute_metrics_multilabel,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=1)]
    )
    
    print("Training Emotion (1 epoch)...")
    trainer.train()
    
    model.save_pretrained("./models/distilroberta-crypto-emotion")
    tokenizer.save_pretrained("./models/distilroberta-crypto-emotion")
    print("✅ Emotion model fine-tuned and saved!")
    return model

# ============================================================
# MAIN
# ============================================================

def main():
    print("="*60)
    print("DOMAIN ADAPTATION: FINE-TUNING WITH LABELED DATA")
    print("="*60)
    
    # Get absolute paths
    base_path = Path("/mnt/dolphinng5_predict/sentiment_engine")
    
    # Load labeled data
    label_file = base_path / "data/labeled_verified.jsonl"
    print(f"\nLoading labeled data from {label_file}...")
    data = load_labeled_data(str(label_file))
    
    sentiment_texts, sentiment_labels = data['sentiment']
    event_texts, event_labels = data['event']
    emotion_texts, emotion_labels = data['emotion']
    
    print(f"Verified samples: {len(sentiment_texts)}")
    
    if len(sentiment_texts) < 5:
        print("⚠️ Not enough verified samples for fine-tuning!")
        return
    
    # Fine-tune sentiment
    fine_tune_sentiment(
        str(base_path / "models/finbert-crypto-sentiment"),
        sentiment_texts, sentiment_labels
    )
    
    # Fine-tune events
    fine_tune_events(
        str(base_path / "models/bert-crypto-events"),
        event_texts, event_labels
    )
    
    # Fine-tune emotion
    fine_tune_emotion(
        str(base_path / "models/distilroberta-crypto-emotion"),
        emotion_texts, emotion_labels
    )
    
    print("\n" + "="*60)
    print("✅ ALL MODELS FINE-TUNED WITH LABELED DATA!")
    print("="*60)

if __name__ == "__main__":
    main()