backend/main.py

#!/usr/bin/env python3
"""EUMORA - Main entry point for emotion analysis."""

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from src.config import config


def main():
    parser = argparse.ArgumentParser(
        description="EUMORA - Emotion-Aware Music Recommendation System",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python main.py train                          # Train on dair-ai/emotion (16k samples)
  python main.py train --goemotions             # Train on GoEmotions (43k samples, better)
  python main.py train --combined               # Best: combine both datasets
  python main.py train --sample                 # Quick test on 2k samples
  python main.py predict "your lyrics"          # Predict emotion
    python main.py predict "your lyrics" --target-sarcasm-prior 0.15
  python main.py predict "text" --detailed-chart # With visualization
  python main.py demo                           # Run demo
  python main.py analyze                        # Interactive mode
        """
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Train command
    train_parser = subparsers.add_parser("train", help="Train the emotion classifier")
    train_parser.add_argument(
        "--sample", action="store_true",
        help="Use 2000 samples for quick training (~5 min)"
    )
    train_parser.add_argument(
        "--samples", type=int, default=None,
        help="Number of training samples to use"
    )
    train_parser.add_argument(
        "--goemotions", action="store_true",
        help="Use GoEmotions dataset (43k samples, 7 emotions)"
    )
    train_parser.add_argument(
        "--combined", action="store_true",
        help="Combine GoEmotions + dair-ai/emotion for best coverage"
    )
    train_parser.add_argument(
        "--no-weights", action="store_true",
        help="Disable class weights (not recommended)"
    )

    # Predict command
    predict_parser = subparsers.add_parser("predict", help="Predict emotion from text")
    predict_parser.add_argument("text", type=str, help="Text to analyze")
    predict_parser.add_argument(
        "--chart", action="store_true",
        help="Force simple bar chart (same as default behavior)"
    )
    predict_parser.add_argument(
        "--detailed-chart", action="store_true",
        help="Generate enhanced detailed analysis chart with primary emotion indicator"
    )
    predict_parser.add_argument(
        "--target-sarcasm-prior", type=float, default=config.target_sarcasm_prior,
        help="Target sarcasm prevalence for lyric-domain prior correction (0-1)"
    )
    predict_parser.add_argument(
        "--train-sarcasm-prior", type=float, default=None,
        help="Optional training-time sarcasm prevalence override (0-1)"
    )
    predict_parser.add_argument(
        "--sarcasm-threshold", type=float, default=config.sarcasm_threshold,
        help="Optional one-vs-rest sarcasm decision threshold (0-1)"
    )
    predict_parser.add_argument(
        "--disable-prior-adjustment", action="store_true",
        help="Disable sarcasm prior-logit adjustment"
    )

    # Demo command
    subparsers.add_parser("demo", help="Run prediction demo")

    # Interactive analyze command
    analyze_parser = subparsers.add_parser("analyze", help="Interactive analysis mode")
    analyze_parser.add_argument(
        "--target-sarcasm-prior", type=float, default=config.target_sarcasm_prior,
        help="Target sarcasm prevalence for lyric-domain prior correction (0-1)"
    )
    analyze_parser.add_argument(
        "--train-sarcasm-prior", type=float, default=None,
        help="Optional training-time sarcasm prevalence override (0-1)"
    )
    analyze_parser.add_argument(
        "--sarcasm-threshold", type=float, default=config.sarcasm_threshold,
        help="Optional one-vs-rest sarcasm decision threshold (0-1)"
    )
    analyze_parser.add_argument(
        "--disable-prior-adjustment", action="store_true",
        help="Disable sarcasm prior-logit adjustment"
    )

    # Recommend command
    recommend_parser = subparsers.add_parser("recommend", help="Recommend Spotify songs from emotional text")
    recommend_parser.add_argument("text", type=str, help="Text describing how you feel")
    recommend_parser.add_argument("--limit", type=int, default=10, help="Number of tracks to return")
    recommend_parser.add_argument("--genre", type=str, default=None, help="Override seed genre (e.g. pop, metal, chill)")
    recommend_parser.add_argument("--no-blend", action="store_true", help="Use top-1 emotion only, no blending")
    recommend_parser.add_argument("--max-popularity", type=int, default=65, help="Max track popularity 0-100 (lower = less mainstream)")
    recommend_parser.add_argument("--target-sarcasm-prior", type=float, default=config.target_sarcasm_prior)
    recommend_parser.add_argument("--disable-prior-adjustment", action="store_true")

    args = parser.parse_args()

    if args.command == "train":
        from src.train import train
        print("\n🎵 EUMORA - Training Emotion Classifier")
        print("=" * 50)

        # Determine dataset
        if args.combined:
            print("📦 Using combined dataset (GoEmotions + dair-ai/emotion)")
        elif args.goemotions:
            print("📦 Using GoEmotions dataset (43k samples, 7 emotions)")
        else:
            print("📦 Using dair-ai/emotion dataset (16k samples, 6 emotions)")

        model_path = train(
            use_sample=args.sample,
            num_train_samples=args.samples,
            use_goemotions=args.goemotions or args.combined,
            combine_datasets=args.combined,
            use_class_weights=not args.no_weights,
        )
        print(f"\n🎉 Training complete!")
        print(f"   Model saved to: {model_path}")
        print(f"\n   Now run: python main.py demo")

    elif args.command == "predict":
        from src.predict import EmotionPredictor
        try:
            target_prior = None if args.disable_prior_adjustment else args.target_sarcasm_prior
            predictor = EmotionPredictor(
                enable_viz=True,
                target_sarcasm_prior=target_prior,
                sarcasm_threshold=args.sarcasm_threshold,
                train_sarcasm_prior=args.train_sarcasm_prior,
            )

            # Determine chart type - always generate a chart
            if args.detailed_chart:
                result = predictor.predict_with_visualization(args.text, chart_type="detailed")
            elif args.chart:
                result = predictor.predict(args.text, create_chart=True, show_chart=True)
            else:
                # Default: generate simple bar chart automatically
                result = predictor.predict(args.text, create_chart=True, show_chart=True)

            print("\n🎵 EUMORA - Emotion Analysis\n")
            print(f"📝 Input: \"{args.text[:80]}{'...' if len(args.text) > 80 else ''}\"")
            print(f"\n🎭 Emotion: {result['emotion'].upper()}")
            print(f"📊 Confidence: {result['confidence']:.1%}")
            print(f"🎸 Music Context: {result['music_context']}")
            print(f"\n💬 {result['explanation']}")

            calibration = result.get('calibration', {})
            print("\n🧪 Sarcasm Calibration:")
            print(f"   target prior: {calibration.get('target_sarcasm_prior')}")
            print(f"   train prior:  {calibration.get('train_sarcasm_prior')}")
            print(f"   logit shift:  {calibration.get('sarcasm_logit_shift')}")
            print(f"   threshold:    {calibration.get('sarcasm_threshold')}")
            print(f"   P(sarcasm):   {calibration.get('sarcasm_probability'):.1%}")

            print("\n📈 All Emotions:")
            sorted_probs = sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True)
            for emotion, prob in sorted_probs:
                bar = "█" * int(prob * 25)
                print(f"   {emotion:>10}: {bar:<25} {prob:.1%}")

            # Show chart info if generated
            if result.get('chart_path'):
                print(f"\n📊 Chart saved to: {result['chart_path']}")

        except FileNotFoundError as e:
            print(f"\n{e}")
            sys.exit(1)

    elif args.command == "recommend":
        try:
            from dotenv import load_dotenv
            load_dotenv()
        except ImportError:
            pass
        from src.predict import EmotionPredictor
        from src.spotify import SpotifyRecommender

        target_prior = None if args.disable_prior_adjustment else args.target_sarcasm_prior

        try:
            predictor   = EmotionPredictor(enable_viz=False, target_sarcasm_prior=target_prior)
            recommender = SpotifyRecommender(max_popularity=args.max_popularity)
        except (ValueError, ImportError, FileNotFoundError) as e:
            print(f"\n\u274c {e}")
            sys.exit(1)

        print(f"\n\U0001f3b5 EUMORA \u2014 Analysing: \"{args.text[:80]}\"")
        predict_result = predictor.predict(args.text)

        print(f"\U0001f3ad Emotion : {predict_result['emotion'].upper()} ({predict_result['confidence']:.1%})")
        print(f"\U0001f3b8 Context : {predict_result['music_context']}")
        print(f"\n\U0001f50d Fetching Spotify recommendations (max popularity: {args.max_popularity})...\n")

        result = recommender.recommend(
            predict_result,
            limit=args.limit,
            genre_override=args.genre,
            blend=not args.no_blend,
        )

        if not result["tracks"]:
            print("\u26a0\ufe0f  No tracks found. Try a different genre or widen the popularity range.")
            sys.exit(0)

        t = result["targets_used"]
        print(
            f"\U0001f3af Targets: valence={t.get('target_valence')}  energy={t.get('target_energy')}  "
            f"dance={t.get('target_danceability')}  tempo={t.get('target_tempo')}  genres={t.get('seed_genres')}\n"
        )

        for i, track in enumerate(result["tracks"], 1):
            af = track["audio_features"]
            print(f"  {i:>2}. {track['name']} \u2014 {track['artist']}")
            print(
                f"      Match: {track['match_score']}%  |  "
                f"val={af['valence']}  nrg={af['energy']}  dance={af['danceability']}  "
                f"tempo={af['tempo']}  pop={track['popularity']}"
            )
            print(f"      {track['spotify_url']}")
            if track.get("preview_url"):
                print(f"      Preview: {track['preview_url']}")
            print()

    elif args.command == "demo":
        from src.predict import demo
        demo()

    elif args.command == "analyze":
        target_prior = None if args.disable_prior_adjustment else args.target_sarcasm_prior
        interactive_mode(
            target_sarcasm_prior=target_prior,
            train_sarcasm_prior=args.train_sarcasm_prior,
            sarcasm_threshold=args.sarcasm_threshold,
        )

    else:
        parser.print_help()


def interactive_mode(
    target_sarcasm_prior=None,
    train_sarcasm_prior=None,
    sarcasm_threshold=None,
):
    """Interactive text analysis mode."""
    print("\n" + "=" * 50)
    print("🎵 EUMORA - Interactive Analysis")
    print("=" * 50)

    try:
        from src.predict import EmotionPredictor
        predictor = EmotionPredictor(
            enable_viz=True,
            target_sarcasm_prior=target_sarcasm_prior,
            train_sarcasm_prior=train_sarcasm_prior,
            sarcasm_threshold=sarcasm_threshold,
        )
    except FileNotFoundError as e:
        print(f"\n{e}")
        return

    print("\nCommands:")
    print("  text                - Analyze text")
    print("  chart: text         - Analyze with simple chart")
    print("  detailed: text      - Analyze with detailed chart")
    print("  quit                - Exit\n")

    while True:
        try:
            user_input = input(">>> ").strip()
            if user_input.lower() in ['quit', 'exit', 'q']:
                print("\n👋 Goodbye!")
                break
            if not user_input:
                continue

            # Parse command
            if user_input.startswith("chart:"):
                text = user_input[6:].strip()
                result = predictor.predict(text, create_chart=True, show_chart=True)
            elif user_input.startswith("detailed:"):
                text = user_input[9:].strip()
                result = predictor.predict_with_visualization(text, chart_type="detailed")
            else:
                text = user_input
                result = predictor.predict(text)

            print(f"    🎭 {result['emotion'].upper()} ({result['confidence']:.1%})")
            print(f"    💬 {result['explanation']}")

            if result.get('chart_path'):
                print(f"    📊 Chart: {result['chart_path']}")
            print()

        except (EOFError, KeyboardInterrupt):
            print("\n👋 Goodbye!")
            break


if __name__ == "__main__":
    main()