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

Domain-specific heads:
  1. FinBERT — financial news sentiment
  2. Financial-RoBERTa-large — policy/formal text sentiment
  3. ClimateBERT — climate risk/opportunity
  4. 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
"""

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


# ─── Default model names ─────────────────────────────────────────
# Fine-tuned on 20K combined financial sentiment corpus (v0.2.0)
DEFAULT_FINBERT = "peyterho/finbert-macro-sentiment"
DEFAULT_ROBERTA = "peyterho/financial-roberta-large-macro-sentiment"
DEFAULT_CLIMATEBERT = "peyterho/climatebert-macro-sentiment"

# Original off-the-shelf models (v0.1.0)
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.
    
    Args:
        model_name: HF model ID. Default: fine-tuned 'peyterho/finbert-macro-sentiment'.
                    Use 'ProsusAI/finbert' for the original off-the-shelf model.
        device: 'cpu' or 'cuda:0'.
    """
    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 sentiment.
    
    Args:
        model_name: HF model ID. Default: fine-tuned 'peyterho/financial-roberta-large-macro-sentiment'.
                    Use 'soleimanian/financial-roberta-large-sentiment' for the original.
        device: 'cpu' or 'cuda:0'.
    """
    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 sentiment.
    
    Args:
        model_name: HF model ID. Default: fine-tuned 'peyterho/climatebert-macro-sentiment'.
                    Use 'climatebert/distilroberta-base-climate-sentiment' for the original.
        device: 'cpu' or 'cuda:0'.
    """
    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 TopicRouter:
    """Keyword-based topic router for macroeconomic text."""
    
    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"):
        pass
    
    def load(self):
        return self
    
    def classify(self, text):
        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]
        
        if best_score < 2:
            return {"topic": "Financial News", "topic_confidence": 0.5, "recommended_head": "finbert", "all_topics": scores}
        elif best_domain == "policy":
            return {"topic": "Macro/Policy", "topic_confidence": min(1.0, policy_hits / 6), "recommended_head": "policy", "all_topics": scores}
        elif best_domain == "climate":
            return {"topic": "Climate/ESG", "topic_confidence": min(1.0, climate_hits / 4), "recommended_head": "climate", "all_topics": scores}
        else:
            return {"topic": "Social/Tweet", "topic_confidence": min(1.0, social_hits / 5), "recommended_head": "finbert", "all_topics": scores}


class TransformerEnsemble:
    """Multi-head transformer ensemble with domain 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.
    
    Default models are fine-tuned on 20K combined financial corpus.
    Pass the ORIGINAL_* constants to use off-the-shelf models.
    
    Examples:
        # Use fine-tuned models (default, recommended)
        ensemble = TransformerEnsemble(device="cpu")
        
        # Use original off-the-shelf models
        from macro_sentiment.transformers_ensemble import ORIGINAL_FINBERT, ORIGINAL_ROBERTA, ORIGINAL_CLIMATEBERT
        ensemble = TransformerEnsemble(
            finbert_model=ORIGINAL_FINBERT,
            roberta_model=ORIGINAL_ROBERTA,
            climatebert_model=ORIGINAL_CLIMATEBERT,
        )
        
        # Mix and match
        ensemble = TransformerEnsemble(
            finbert_model="peyterho/finbert-macro-sentiment",
            roberta_model="soleimanian/financial-roberta-large-sentiment",
        )
    """
    def __init__(
        self,
        device="cpu",
        use_router=True,
        finbert_model=DEFAULT_FINBERT,
        roberta_model=DEFAULT_ROBERTA,
        climatebert_model=DEFAULT_CLIMATEBERT,
    ):
        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),
        }
        self.router = TopicRouter(device) 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"
        sentiment = self.heads[head_name].score(text)
        return {
            "head_used": head_name,
            "topic": topic_info["topic"],
            "topic_confidence": topic_info["topic_confidence"],
            "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