#!/usr/bin/env python3
"""
AENEA BHR-0.1 Inference Engine
================================
Production-ready inference wrapper for the BHR decision model.

Usage:
    from bhr_engine import BHREngine
    
    engine = BHREngine("bhr_best.pt", "tokenizers/QT-Cittern-12k")
    
    result = engine.decide(
        state='text_A: The cat sat on the mat.\\ntext_B: The mat was empty.',
        question='Does text_A contradict text_B? Is "contradiction" the correct answer?',
        options=["yes", "no"]
    )
    
    print(result.decision)      # "yes"
    print(result.confidence)    # 0.93
    print(result.probabilities) # {"yes": 0.93, "no": 0.07}
"""

import time
import torch
import torch.nn.functional as F
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Union
from transformers import AutoTokenizer

from cartan_olm import CartanConfig, CartanLM


# ═══════════════════════════════════════════════════════════════════════════
# DATA STRUCTURES
# ═══════════════════════════════════════════════════════════════════════════

@dataclass
class DecisionResult:
    """Result of a single decision."""
    decision: str                    # The chosen option
    confidence: float                # Probability of the chosen option
    probabilities: Dict[str, float]  # Full probability distribution
    latency_ms: float                # Inference time in milliseconds
    kind: str                        # "choice", "yesno", or "score"
    options: List[str]               # All options presented
    
    def __repr__(self):
        return (f"DecisionResult(decision='{self.decision}', "
                f"confidence={self.confidence:.4f}, "
                f"latency={self.latency_ms:.1f}ms)")


@dataclass
class BatchResult:
    """Result of a batch decision."""
    results: List[DecisionResult]
    total_latency_ms: float
    
    def __repr__(self):
        return (f"BatchResult(n={len(self.results)}, "
                f"total_latency={self.total_latency_ms:.1f}ms)")


# ═══════════════════════════════════════════════════════════════════════════
# ENGINE
# ═══════════════════════════════════════════════════════════════════════════

class BHREngine:
    """
    Production inference engine for BHR-0.1 decision model.
    
    Args:
        checkpoint_path: Path to the BHR checkpoint (.pt file)
        tokenizer_path: Path to the tokenizer directory
        device: "cpu" or "cuda"
        max_seq_len: Maximum sequence length for state+question
        max_option_len: Maximum token length per option
    """
    
    def __init__(
        self,
        checkpoint_path: str = "bhr_best.pt",
        tokenizer_path: str = "tokenizers/QT-Cittern-12k",
        device: str = "cpu",
        max_seq_len: int = 1536,
        max_option_len: int = 64,
    ):
        self.device = device
        self.max_seq_len = max_seq_len
        self.max_option_len = max_option_len
        
        # Load checkpoint
        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
        
        # Build config
        self.config = CartanConfig.from_dict(checkpoint["config"])
        self.config.decision_mode = True
        
        # Load tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.config.vocab_size = len(self.tokenizer)
        
        # Build and load model
        self.model = CartanLM(self.config).to(device)
        state_dict = {
            (k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k): v
            for k, v in checkpoint["model_state_dict"].items()
        }
        self.model.load_state_dict(state_dict, strict=False)
        self.model.eval()
        
        # Store metadata
        self.checkpoint_metadata = {
            "step": checkpoint.get("step"),
            "epoch": checkpoint.get("epoch"),
            "test_accuracy": checkpoint.get("test_accuracy"),
            "test_brier": checkpoint.get("test_brier"),
            "test_ece": checkpoint.get("test_ece"),
        }
        
        self.pad_id = self.tokenizer.pad_token_id or 0
        self._param_count = sum(p.numel() for p in self.model.parameters())
        
        print(f"BHR Engine initialized: {self._param_count / 1e6:.1f}M params on {device}")
    
    @property
    def param_count_m(self) -> float:
        """Parameter count in millions."""
        return self._param_count / 1e6
    
    def _encode_state_question(self, state: str, question: str) -> torch.Tensor:
        """Encode state + question into input_ids tensor."""
        state_text = f"State: {state}\nQuestion: {question}"
        ids = self.tokenizer.encode(
            state_text,
            add_special_tokens=False,
            truncation=True,
            max_length=self.max_seq_len
        )
        return torch.tensor([ids], dtype=torch.long, device=self.device)
    
    def _encode_options(self, options: List[str]) -> torch.Tensor:
        """Encode and pad options into option_ids tensor."""
        token_lists = []
        for opt in options:
            ids = self.tokenizer.encode(
                str(opt),
                add_special_tokens=False,
                truncation=True,
                max_length=self.max_option_len
            )
            if not ids:
                ids = [self.pad_id]
            token_lists.append(ids)
        
        max_len = max(len(t) for t in token_lists)
        padded = []
        for tokens in token_lists:
            padded.append(tokens + [self.pad_id] * (max_len - len(tokens)))
        
        return torch.tensor([padded], dtype=torch.long, device=self.device)
    
    def _detect_kind(self, options: List[str]) -> str:
        """Detect the decision kind from options."""
        if len(options) == 2:
            lowered = [o.lower() for o in options]
            if set(lowered) == {"yes", "no"}:
                return "yesno"
            if set(lowered) == {"true", "false"}:
                return "yesno"
        # Check if options are numeric or ordinal
        try:
            [float(o) for o in options]
            return "score"
        except ValueError:
            pass
        if any("level" in o.lower() for o in options):
            return "score"
        return "choice"
    
    def decide(
        self,
        state: str,
        question: str,
        options: List[str],
        return_logits: bool = False,
    ) -> DecisionResult:
        """
        Make a single decision.
        
        Args:
            state: The context/state to evaluate
            question: The question to answer
            options: List of possible answers
            return_logits: If True, include raw logits in result
            
        Returns:
            DecisionResult with decision, confidence, and probabilities
        """
        if not options:
            raise ValueError("options must not be empty")
        if len(options) == 1:
            return DecisionResult(
                decision=options[0],
                confidence=1.0,
                probabilities={options[0]: 1.0},
                latency_ms=0.0,
                kind=self._detect_kind(options),
                options=options,
            )
        
        input_ids = self._encode_state_question(state, question)
        option_ids = self._encode_options(options)
        
        start = time.perf_counter()
        with torch.no_grad():
            output = self.model(
                input_ids,
                option_ids=option_ids,
                pad_token_id=self.pad_id
            )
        latency = (time.perf_counter() - start) * 1000
        
        logits = output["logits"][0].float()
        probs = F.softmax(logits, dim=-1).cpu().tolist()
        best_idx = int(torch.argmax(logits).item())
        
        prob_dict = {opt: p for opt, p in zip(options, probs)}
        
        result = DecisionResult(
            decision=options[best_idx],
            confidence=probs[best_idx],
            probabilities=prob_dict,
            latency_ms=latency,
            kind=self._detect_kind(options),
            options=options,
        )
        
        return result
    
    def batch_decide(
        self,
        items: List[Dict[str, Union[str, List[str]]]],
    ) -> BatchResult:
        """
        Make decisions on a batch of items.
        
        Each item should have:
            - "state": str
            - "question": str
            - "options": List[str]
            
        Returns:
            BatchResult with list of DecisionResults
        """
        results = []
        start = time.perf_counter()
        
        for item in items:
            result = self.decide(
                state=item["state"],
                question=item["question"],
                options=item["options"],
            )
            results.append(result)
        
        total_latency = (time.perf_counter() - start) * 1000
        return BatchResult(results=results, total_latency_ms=total_latency)
    
    def yesno(self, state: str, question: str) -> DecisionResult:
        """Convenience method for yes/no decisions."""
        return self.decide(state, question, ["yes", "no"])
    
    def choose(self, state: str, question: str, options: List[str]) -> DecisionResult:
        """Convenience method for choice decisions."""
        return self.decide(state, question, options)
    
    def score(self, state: str, question: str, options: List[str]) -> DecisionResult:
        """Convenience method for score decisions."""
        return self.decide(state, question, options)
    
    def info(self) -> Dict:
        """Return engine metadata."""
        return {
            "model": "AENEA BHR-0.1",
            "parameters_m": round(self.param_count_m, 1),
            "architecture": "CartanLM (Orthogonal)",
            "device": self.device,
            "max_seq_len": self.max_seq_len,
            "max_option_len": self.max_option_len,
            **self.checkpoint_metadata,
        }


# ═══════════════════════════════════════════════════════════════════════════
# CLI INTERFACE
# ═══════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="BHR-0.1 Inference Engine")
    parser.add_argument("--checkpoint", default="bhr_best.pt")
    parser.add_argument("--tokenizer", default="tokenizers/QT-Cittern-12k")
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--interactive", action="store_true")
    args = parser.parse_args()
    
    engine = BHREngine(args.checkpoint, args.tokenizer, args.device)
    print(f"\nEngine info: {engine.info()}\n")
    
    if args.interactive:
        print("BHR Interactive Mode (type 'quit' to exit)")
        print("=" * 60)
        while True:
            state = input("\nState: ").strip()
            if state.lower() == "quit":
                break
            question = input("Question: ").strip()
            if question.lower() == "quit":
                break
            options_str = input("Options (comma-separated): ").strip()
            if options_str.lower() == "quit":
                break
            options = [o.strip() for o in options_str.split(",") if o.strip()]
            if not options:
                continue
            
            result = engine.decide(state, question, options)
            print(f"\n  Decision: {result.decision}")
            print(f"  Confidence: {result.confidence:.4f}")
            print(f"  Kind: {result.kind}")
            print(f"  Latency: {result.latency_ms:.1f}ms")
            print(f"  Distribution:")
            for opt, prob in sorted(result.probabilities.items(), key=lambda x: -x[1]):
                bar = "█" * int(prob * 30)
                print(f"    {prob:6.4f} {bar:<30s} {opt}")
    else:
        # Run a quick demo
        result = engine.yesno(
            state='{"user": {"age": 25}, "rule": "Allow entry if age >= 18."}',
            question="Following the rule, is entry allowed?"
        )
        print(f"Demo result: {result}")