"""
Meta-classifier: trains a GradientBoosting model on dictionary + (optionally) transformer features
to predict 3-class sentiment labels (negative/neutral/positive).

Two modes:
  1. Dict-only: 24 dictionary features → fast baseline (F1=0.58)
  2. Combined: 24 dict + transformer ensemble features → best accuracy

Usage:
  python -m macro_sentiment.train_meta [--use-transformers] [--output PATH]
"""
import os
import sys
import time
import json
import pickle
import numpy as np
from collections import Counter
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import classification_report, f1_score, accuracy_score
from sklearn.model_selection import cross_val_score

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from macro_sentiment.dictionaries import CombinedDictionaryScorer


def load_data():
    """Load combined dataset."""
    from macro_sentiment.data_prep import load_combined_dataset
    ds = load_combined_dataset()
    return ds["train"], ds["test"]


def extract_dict_features(texts, scorer):
    """Extract dictionary features for a list of texts."""
    features = []
    for i, text in enumerate(texts):
        if i % 500 == 0:
            print(f"  Scoring {i}/{len(texts)}...")
            sys.stdout.flush()
        feat = scorer.score(text)
        features.append(feat)
    
    keys = sorted(features[0].keys())
    X = np.array([[f[k] for k in keys] for f in features])
    return X, keys


def train_meta_classifier(use_transformers=False, output_path="meta_classifier.pkl"):
    print(f"[{time.strftime('%H:%M:%S')}] Training meta-classifier (transformers={use_transformers})")
    
    train_ds, test_ds = load_data()
    train_texts = train_ds["text"]
    test_texts = test_ds["text"]
    y_train = np.array(train_ds["label"])
    y_test = np.array(test_ds["label"])
    
    print(f"\n[{time.strftime('%H:%M:%S')}] Extracting dictionary features...")
    sys.stdout.flush()
    
    scorer = CombinedDictionaryScorer()
    X_train_dict, feature_names = extract_dict_features(train_texts, scorer)
    X_test_dict, _ = extract_dict_features(test_texts, scorer)
    
    X_train = X_train_dict
    X_test = X_test_dict
    all_feature_names = feature_names
    
    if use_transformers:
        print(f"\n[{time.strftime('%H:%M:%S')}] Extracting transformer features...")
        sys.stdout.flush()
        from macro_sentiment.transformers_ensemble import TransformerEnsemble
        ensemble = TransformerEnsemble(device="cpu", use_router=True)
        ensemble.load_all()
        
        tf_features_train = []
        for i, text in enumerate(train_texts):
            if i % 100 == 0:
                print(f"  Transformer scoring {i}/{len(train_texts)}...")
                sys.stdout.flush()
            try:
                scores = ensemble.score_all(text)
                tf_features_train.append(scores)
            except Exception as e:
                print(f"  Error on sample {i}: {e}")
                tf_features_train.append({})
        
        tf_keys = sorted(set().union(*(f.keys() for f in tf_features_train if f)))
        tf_keys = [k for k in tf_keys if isinstance(tf_features_train[0].get(k, 0), (int, float))]
        
        X_train_tf = np.array([[f.get(k, 0.0) for k in tf_keys] for f in tf_features_train])
        X_train = np.hstack([X_train_dict, X_train_tf])
        
        tf_features_test = []
        for i, text in enumerate(test_texts):
            if i % 100 == 0:
                print(f"  Transformer scoring test {i}/{len(test_texts)}...")
                sys.stdout.flush()
            try:
                scores = ensemble.score_all(text)
                tf_features_test.append(scores)
            except Exception as e:
                tf_features_test.append({})
        
        X_test_tf = np.array([[f.get(k, 0.0) for k in tf_keys] for f in tf_features_test])
        X_test = np.hstack([X_test_dict, X_test_tf])
        all_feature_names = feature_names + [f"tf_{k}" for k in tf_keys]
    
    print(f"\n[{time.strftime('%H:%M:%S')}] Training GradientBoosting...")
    print(f"  Features: {X_train.shape[1]}")
    print(f"  Train samples: {X_train.shape[0]}")
    sys.stdout.flush()
    
    clf = GradientBoostingClassifier(
        n_estimators=200,
        max_depth=5,
        learning_rate=0.1,
        subsample=0.8,
        random_state=42,
    )
    clf.fit(X_train, y_train)
    
    y_pred = clf.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred, average="macro")
    report = classification_report(y_test, y_pred, target_names=["negative", "neutral", "positive"])
    
    print(f"\n[{time.strftime('%H:%M:%S')}] Results:")
    print(f"  Accuracy: {acc:.4f}")
    print(f"  F1 macro: {f1:.4f}")
    print(f"\n{report}")
    sys.stdout.flush()
    
    importances = sorted(zip(all_feature_names, clf.feature_importances_), key=lambda x: -x[1])
    print("\nTop 10 Features:")
    for name, imp in importances[:10]:
        print(f"  {name}: {imp:.4f}")
    
    artifact = {
        "model": clf,
        "feature_names": all_feature_names,
        "accuracy": acc,
        "f1_macro": f1,
        "use_transformers": use_transformers,
        "report": report,
        "feature_importances": importances,
    }
    with open(output_path, "wb") as f:
        pickle.dump(artifact, f)
    print(f"\n[{time.strftime('%H:%M:%S')}] Saved to {output_path}")
    
    return artifact


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--use-transformers", action="store_true")
    parser.add_argument("--output", default="meta_classifier.pkl")
    args = parser.parse_args()
    train_meta_classifier(use_transformers=args.use_transformers, output_path=args.output)