"""
Multi-head transformer ensemble for macroeconomic sentiment.

Domain-specific heads:
  1. FinBERT — financial news sentiment (English)
  2. Financial-RoBERTa-large — policy/formal text sentiment (English)
  3. ClimateBERT — climate risk/opportunity (English)
  4. XLM-RoBERTa — multilingual sentiment (8+ languages)
  5. Keyword-based Topic Router — routes to appropriate head

All heads output standardized scores in [-1, +1].

Fine-tuned models (v0.2.0):
  - FinBERT:    peyterho/finbert-macro-sentiment
  - RoBERTa:    peyterho/financial-roberta-large-macro-sentiment
  - ClimateBERT: peyterho/climatebert-macro-sentiment

Multilingual (v0.3.0):
  - XLM-RoBERTa: cardiffnlp/twitter-xlm-roberta-base-sentiment-multilingual
"""

import re
import numpy as np
from typing import Dict, List, Optional, Tuple
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    pipeline,
)


# ─── Default model names ─────────────────────────────────────────
DEFAULT_FINBERT = "peyterho/finbert-macro-sentiment"
DEFAULT_ROBERTA = "peyterho/financial-roberta-large-macro-sentiment"
DEFAULT_CLIMATEBERT = "peyterho/climatebert-macro-sentiment"
DEFAULT_MULTILINGUAL = "cardiffnlp/twitter-xlm-roberta-base-sentiment-multilingual"

ORIGINAL_FINBERT = "ProsusAI/finbert"
ORIGINAL_ROBERTA = "soleimanian/financial-roberta-large-sentiment"
ORIGINAL_CLIMATEBERT = "climatebert/distilroberta-base-climate-sentiment"


class SentimentHead:
    """Base class for a transformer sentiment scoring head."""

    def __init__(self, model_name, label_map, score_map, device="cpu", max_length=512):
        self.model_name = model_name
        self.label_map = label_map
        self.score_map = score_map
        self.device = device
        self.max_length = max_length
        self._pipeline = None

    def load(self):
        if self._pipeline is None:
            tokenizer = AutoTokenizer.from_pretrained(self.model_name, model_max_length=self.max_length)
            model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
            self._pipeline = pipeline(
                "text-classification", model=model, tokenizer=tokenizer,
                device=self.device if self.device == "cpu" else int(self.device.split(":")[-1]),
                truncation=True, max_length=self.max_length, top_k=None,
            )
        return self

    def score(self, text):
        self.load()
        outputs = self._pipeline(text)
        if outputs and isinstance(outputs[0], list):
            outputs = outputs[0]

        result = {}
        weighted_score = 0.0
        for item in outputs:
            label = item["label"]
            prob = item["score"]
            result[f"prob_{label}"] = prob
            if label in self.score_map:
                weighted_score += prob * self.score_map[label]
            else:
                for idx, name in self.label_map.items():
                    if label == f"LABEL_{idx}" or label == name:
                        weighted_score += prob * self.score_map.get(name, 0.0)
                        result[f"prob_{name}"] = prob
                        break
        result["composite_score"] = weighted_score
        return result


class FinBERTHead(SentimentHead):
    """FinBERT for financial news sentiment (English)."""
    def __init__(self, model_name=DEFAULT_FINBERT, device="cpu"):
        super().__init__(
            model_name=model_name,
            label_map={0: "positive", 1: "negative", 2: "neutral"},
            score_map={"positive": 1.0, "negative": -1.0, "neutral": 0.0},
            device=device,
        )


class FinancialRoBERTaHead(SentimentHead):
    """Financial-RoBERTa-Large for policy/formal text (English)."""
    def __init__(self, model_name=DEFAULT_ROBERTA, device="cpu"):
        super().__init__(
            model_name=model_name,
            label_map={0: "negative", 1: "neutral", 2: "positive"},
            score_map={"positive": 1.0, "negative": -1.0, "neutral": 0.0},
            device=device, max_length=512,
        )


class ClimateBERTHead(SentimentHead):
    """ClimateBERT for climate risk/opportunity (English)."""
    def __init__(self, model_name=DEFAULT_CLIMATEBERT, device="cpu"):
        super().__init__(
            model_name=model_name,
            label_map={0: "risk", 1: "neutral", 2: "opportunity"},
            score_map={"risk": -1.0, "neutral": 0.0, "opportunity": 1.0},
            device=device, max_length=512,
        )


class MultilingualHead(SentimentHead):
    """XLM-RoBERTa for multilingual sentiment (8+ languages).

    Supports: English, Arabic, French, German, Hindi, Italian, Portuguese, Spanish
    and performs reasonably on many other languages.

    Uses cardiffnlp/twitter-xlm-roberta-base-sentiment-multilingual by default.

    Args:
        model_name: HF model ID for the multilingual model.
        device: 'cpu' or 'cuda:0'.
    """
    def __init__(self, model_name=DEFAULT_MULTILINGUAL, device="cpu"):
        super().__init__(
            model_name=model_name,
            label_map={0: "negative", 1: "neutral", 2: "positive"},
            score_map={"positive": 1.0, "negative": -1.0, "neutral": 0.0},
            device=device, max_length=512,
        )


# ─── Language Detection ──────────────────────────────────────────

def detect_language(text):
    """Simple heuristic language detection. Returns 'en' or 'other'.

    Uses Unicode script analysis — no external dependencies.
    For production use, install langdetect or fasttext for better accuracy.
    """
    # Try to import langdetect if available
    try:
        from langdetect import detect
        lang = detect(text)
        return lang
    except ImportError:
        pass

    # Fallback: heuristic based on character ranges
    ascii_chars = sum(1 for c in text if ord(c) < 128)
    total_chars = max(len(text), 1)

    # If >85% ASCII, likely English (or another Latin-script language)
    # Check for common non-English Latin indicators
    text_lower = text.lower()

    # German indicators
    if any(c in text for c in "äöüß") or any(w in text_lower for w in ["und", "der", "die", "das"]):
        return "de"
    # French indicators
    if any(c in text for c in "éèêëàâùûçœ") or any(w in text_lower for w in [" le ", " la ", " les ", " des "]):
        return "fr"
    # Spanish indicators
    if any(c in text for c in "ñáéíóú¿¡") or any(w in text_lower for w in [" el ", " los ", " las "]):
        return "es"
    # Portuguese
    if any(c in text for c in "ãõç") or any(w in text_lower for w in [" os ", " das ", " nos "]):
        return "pt"
    # Japanese/Chinese (CJK characters)
    if any('\u4e00' <= c <= '\u9fff' or '\u3040' <= c <= '\u30ff' for c in text):
        return "ja"  # could be zh too
    # Arabic
    if any('\u0600' <= c <= '\u06ff' for c in text):
        return "ar"
    # Hindi/Devanagari
    if any('\u0900' <= c <= '\u097f' for c in text):
        return "hi"

    if ascii_chars / total_chars > 0.85:
        return "en"

    return "other"


# ─── Topic Router ────────────────────────────────────────────────

class TopicRouter:
    """Keyword-based topic router with language-aware multilingual fallback.

    Routing logic:
      1. Detect language — if non-English, route to multilingual head
      2. For English text, use keyword matching:
         - Policy/macro keywords → RoBERTa-Large (policy head)
         - Climate/ESG keywords → ClimateBERT
         - Social indicators ($cashtags, slang) → FinBERT
         - Default → FinBERT
    """

    POLICY_KEYWORDS = {
        "federal reserve", "fed ", "fomc", "central bank", "ecb", "boe",
        "bank of japan", "boj", "bank of england", "rba", "pboc",
        "bank of canada", "riksbank", "snb", "reserve bank",
        "lagarde", "powell", "yellen", "draghi", "bernanke",
        "interest rate", "rate hike", "rate cut", "monetary policy",
        "quantitative easing", "quantitative tightening", "tapering",
        "policy rate", "discount rate", "fed funds", "federal funds",
        "basis points", "bps", "hawkish", "dovish", "accommodative",
        "tightening", "easing", "forward guidance", "yield curve",
        "rates unchanged", "data-dependent", "data dependent",
        "policy stance", "policy decision", "rate decision",
        "fiscal policy", "government spending", "treasury", "sovereign",
        "debt ceiling", "fiscal stimulus", "budget deficit", "gdp",
        "unemployment", "labor market", "jobs report", "nonfarm",
        "payrolls", "consumer price index", "cpi", "pce", "inflation",
        "deflation", "stagflation", "ppi", "economic growth",
        "recession", "economic outlook", "macro", "macroeconomic",
    }

    CLIMATE_KEYWORDS = {
        "climate", "carbon", "emission", "emissions", "renewable",
        "sustainability", "sustainable", "esg", "green bond",
        "net zero", "net-zero", "paris agreement", "clean energy",
        "fossil fuel", "decarbonization", "decarbonisation",
        "tcfd", "sfdr", "eu taxonomy", "climate risk",
        "climate change", "global warming", "greenhouse",
        "environmental", "carbon tax", "carbon price",
        "wind energy", "solar energy", "electric vehicle",
    }

    SOCIAL_INDICATORS = {
        "$", "#", "@", "imo", "imho", "lol", "lmao", "tbh",
        "bullish af", "bearish af", "moon", "to the moon",
        "diamond hands", "paper hands", "hodl", "fomo", "yolo",
    }

    def __init__(self, device="cpu", enable_multilingual=True):
        self.enable_multilingual = enable_multilingual

    def load(self):
        return self

    def classify(self, text):
        # Step 1: Language detection
        lang = detect_language(text) if self.enable_multilingual else "en"

        if lang != "en" and self.enable_multilingual:
            return {
                "topic": f"Multilingual ({lang})",
                "topic_confidence": 0.8,
                "recommended_head": "multilingual",
                "detected_language": lang,
                "all_topics": {"multilingual": 10},
            }

        # Step 2: English keyword routing
        text_lower = text.lower()
        policy_hits = sum(1 for kw in self.POLICY_KEYWORDS if kw in text_lower)
        climate_hits = sum(1 for kw in self.CLIMATE_KEYWORDS if kw in text_lower)
        social_hits = sum(1 for ind in self.SOCIAL_INDICATORS if ind in text)

        has_cashtag = bool(re.search(r'\$[A-Z]{1,5}\b', text))
        is_short = len(text.split()) < 40
        if has_cashtag: social_hits += 3
        if is_short: social_hits += 1

        scores = {"policy": policy_hits, "climate": climate_hits, "social": social_hits}
        best_domain = max(scores, key=scores.get)
        best_score = scores[best_domain]

        result = {"detected_language": lang, "all_topics": scores}

        if best_score < 2:
            result.update({"topic": "Financial News", "topic_confidence": 0.5, "recommended_head": "finbert"})
        elif best_domain == "policy":
            result.update({"topic": "Macro/Policy", "topic_confidence": min(1.0, policy_hits / 6), "recommended_head": "policy"})
        elif best_domain == "climate":
            result.update({"topic": "Climate/ESG", "topic_confidence": min(1.0, climate_hits / 4), "recommended_head": "climate"})
        else:
            result.update({"topic": "Social/Tweet", "topic_confidence": min(1.0, social_hits / 5), "recommended_head": "finbert"})

        return result


class TransformerEnsemble:
    """Multi-head transformer ensemble with domain and language routing.

    Args:
        device: 'cpu' or 'cuda:0'.
        use_router: Enable keyword-based topic routing.
        finbert_model: HF model ID for the FinBERT head.
        roberta_model: HF model ID for the Financial-RoBERTa head.
        climatebert_model: HF model ID for the ClimateBERT head.
        multilingual_model: HF model ID for the multilingual head. Set to None to disable.
        enable_multilingual: Auto-route non-English text to multilingual head.

    Examples:
        # All heads including multilingual (default)
        ensemble = TransformerEnsemble(device="cpu")

        # English-only (no multilingual head, v0.2.0 behavior)
        ensemble = TransformerEnsemble(multilingual_model=None)

        # Score German text — auto-routed to multilingual head
        result = ensemble.score_routed("EZB signalisiert Geduld bei Zinssenkungen.")
        # result["head_used"] = "multilingual"
    """
    def __init__(
        self,
        device="cpu",
        use_router=True,
        finbert_model=DEFAULT_FINBERT,
        roberta_model=DEFAULT_ROBERTA,
        climatebert_model=DEFAULT_CLIMATEBERT,
        multilingual_model=DEFAULT_MULTILINGUAL,
        enable_multilingual=True,
    ):
        self.device = device
        self.use_router = use_router
        self.heads = {
            "finbert": FinBERTHead(model_name=finbert_model, device=device),
            "policy": FinancialRoBERTaHead(model_name=roberta_model, device=device),
            "climate": ClimateBERTHead(model_name=climatebert_model, device=device),
        }
        if multilingual_model:
            self.heads["multilingual"] = MultilingualHead(model_name=multilingual_model, device=device)
        self.router = TopicRouter(device, enable_multilingual=enable_multilingual and multilingual_model is not None) if use_router else None

    def score_routed(self, text):
        if not self.router:
            raise ValueError("Router not enabled.")
        topic_info = self.router.classify(text)
        head_name = topic_info["recommended_head"]
        if head_name == "tweet": head_name = "finbert"
        # Fallback if multilingual head not loaded
        if head_name == "multilingual" and "multilingual" not in self.heads:
            head_name = "finbert"
        sentiment = self.heads[head_name].score(text)
        return {
            "head_used": head_name,
            "topic": topic_info["topic"],
            "topic_confidence": topic_info["topic_confidence"],
            "detected_language": topic_info.get("detected_language", "en"),
            "sentiment_score": sentiment["composite_score"],
            **{f"{head_name}_{k}": v for k, v in sentiment.items()},
        }

    def score_all(self, text):
        result = {}
        for name, head in self.heads.items():
            for k, v in head.score(text).items():
                result[f"{name}_{k}"] = v
        composites = [result.get(f"{name}_composite_score", 0.0) for name in self.heads]
        result["ensemble_mean"] = np.mean(composites)
        return result

    def load_all(self):
        for head in self.heads.values():
            head.load()
        if self.router: self.router.load()
        return self