import os
import pickle
import pandas as pd
import xgboost as xgb
from dataclasses import dataclass
from typing import List, Dict, Tuple, Optional
import math
from .base_calculator import BaseCalculator, CalculationContext
from .confidence import calc_confidence_3way, calc_confidence_dc
from .match_result_calculator import MatchResultPrediction

@dataclass
class ScorePrediction:
    predicted_ft_score: str
    predicted_ht_score: str
    ft_scores_top5: List[Dict]
    
    # Reconciled MS/DC predictions (can be updated here)
    reconciled_ms: Optional[MatchResultPrediction] = None

class ScoreCalculator(BaseCalculator):

    def __init__(self, config: Dict):
        super().__init__(config)
        self.xgb_home = None
        self.xgb_away = None
        self.xgb_ht_home = None
        self.xgb_ht_away = None
        self.scaler = None # If used
        self.features = []
        self._load_model()
        
    def _load_model(self):
        try:
            model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "models", "xgb_score.pkl")
            if os.path.exists(model_path):
                with open(model_path, "rb") as f:
                    data = pickle.load(f)
                    # Handle both dictionary and direct model formats (just in case)
                    if isinstance(data, dict):
                        self.xgb_home = data.get("home_model")
                        self.xgb_away = data.get("away_model")
                        self.xgb_ht_home = data.get("ht_home_model")
                        self.xgb_ht_away = data.get("ht_away_model")
                        self.features = data.get("features", [])
                    else:
                        print("⚠️ Unexpected XGB score model format.")
                print("✅ XGBoost Score Model loaded.")
            else:
                print(f"⚠️ XGBoost Score Model not found at {model_path}")
        except Exception as e:
            print(f"❌ Error loading XGBoost Score Model: {e}")

    def _poisson_pmf(self, k, lam):
        """Poisson probability mass function."""
        if lam <= 0:
            return 1.0 if k == 0 else 0.0
        return (lam ** k) * math.exp(-lam) / math.factorial(k)
        
    def calculate(self, ctx: CalculationContext, ms_result: MatchResultPrediction) -> ScorePrediction:  # type: ignore[override]
        predicted_ht = None
        # Default Lambdas (fallback)
        lambda_home = max(0.5, ctx.home_xg)
        lambda_away = max(0.5, ctx.away_xg)
        
        # --- XGBOOST PREDICTION ---
        if self.xgb_home and self.xgb_away and hasattr(ctx.team_pred, "raw_features"):
            try:
                # 1. Prepare Features
                # We need to map ctx data to self.features list columns
                raw = ctx.team_pred.raw_features
                odds = ctx.odds_data or {}
                
                # Use unified feature adapter for exact 56-feature sync
                from features.feature_adapter import get_feature_adapter
                df_input = get_feature_adapter().get_features(ctx)
                
                # Predict FT
                pred_h = self.xgb_home.predict(df_input)[0]
                pred_a = self.xgb_away.predict(df_input)[0]
                
                # Predict HT (if available)
                if self.xgb_ht_home and self.xgb_ht_away:
                    pred_ht_h = self.xgb_ht_home.predict(df_input)[0]
                    pred_ht_a = self.xgb_ht_away.predict(df_input)[0]
                    
                    # Clamp HT predictions (min 0, and shouldn't exceed FT in logic, but models are independent)
                    # We trust the model but ensure sanity (HT <= FT is hard to enforce without joint training, but usually holds)
                    ht_h_val = max(0.0, float(pred_ht_h))
                    ht_a_val = max(0.0, float(pred_ht_a))
                    
                    predicted_ht = f"{round(ht_h_val)}-{round(ht_a_val)}"
                else:
                    # Fallback if HT models missing
                    ht_h_val = max(0.0, float(pred_h) * 0.42)
                    ht_a_val = max(0.0, float(pred_a) * 0.42)
                    predicted_ht = f"{round(ht_h_val)}-{round(ht_a_val)}"
                
                # Update lambdas with ML predictions
                lambda_home = max(0.1, min(6.0, float(pred_h)))
                lambda_away = max(0.1, min(6.0, float(pred_a)))
                
                # Store raw XGB preds in context
                ctx.xgboost_preds["score"] = {
                    "home": lambda_home, 
                    "away": lambda_away,
                    "ht_home": ht_h_val,
                    "ht_away": ht_a_val
                }
                
            except Exception as e:
                print(f"⚠️ XGBoost Score Prediction failed: {e}. Falling back to Poisson xG.")
                # Fallback to current simple logic if ML fails
                predicted_ht = f"{round(lambda_home * 0.42)}-{round(lambda_away * 0.42)}"
        
        # --- POISSON GRID GENERATION ---
        # Now use lambda_home/away (either ML or fallback) to generate grid
        score_probs = {}
        grid_max = self.config.get("score.poisson_grid_max", 7)
        
        for i in range(grid_max):
            for j in range(grid_max):
                p = self._poisson_pmf(i, lambda_home) * self._poisson_pmf(j, lambda_away)
                score_probs[f"{i}-{j}"] = round(p * 100, 2)
        
        sorted_scores = sorted(score_probs.items(), key=lambda x: x[1], reverse=True)
        
        # --- DERIVE MS PROBS FROM SCORES (CONSISTENCY CHECK) ---
        poisson_ms_home = sum(p for s, p in score_probs.items() 
                              for h, a in [s.split("-")] if int(h) > int(a))
        poisson_ms_away = sum(p for s, p in score_probs.items()
                              for h, a in [s.split("-")] if int(h) < int(a))
        poisson_ms_draw = sum(p for s, p in score_probs.items()
                              for h, a in [s.split("-")] if int(h) == int(a))
        
        # Normalize
        poisson_total = poisson_ms_home + poisson_ms_away + poisson_ms_draw
        if poisson_total > 0:
            poisson_ms_home /= poisson_total
            poisson_ms_away /= poisson_total
            poisson_ms_draw /= poisson_total
            
        # --- HYBRID RECONCILIATION ---
        
        threshold = self.config.get("score.ms_confidence_threshold", 15.0)
        reconciled_result = ms_result
        
        # If original confidence is low, trust new Score Model more
        if ms_result.ms_confidence < threshold:
            poisson_probs = [(poisson_ms_home, "1"), (poisson_ms_draw, "X"), (poisson_ms_away, "2")]
            poisson_sorted = sorted(poisson_probs, key=lambda x: x[0], reverse=True)
            
            new_ms_pick = poisson_sorted[0][1]
            new_ms_conf = calc_confidence_3way(poisson_sorted[0][0])
            
            # Recalculate DC
            dc_1x = poisson_ms_home + poisson_ms_draw
            dc_x2 = poisson_ms_draw + poisson_ms_away
            dc_12 = poisson_ms_home + poisson_ms_away
            
            dc_probs = [(dc_1x, "1X"), (dc_x2, "X2"), (dc_12, "12")]
            dc_sorted = sorted(dc_probs, key=lambda x: x[0], reverse=True)
            new_dc_pick = dc_sorted[0][1]
            new_dc_conf = calc_confidence_dc(dc_sorted[0][0])
            
            reconciled_result = MatchResultPrediction(
                ms_home_prob=poisson_ms_home,
                ms_draw_prob=poisson_ms_draw,
                ms_away_prob=poisson_ms_away,
                ms_pick=new_ms_pick,
                ms_confidence=new_ms_conf,
                dc_1x_prob=dc_1x,
                dc_x2_prob=dc_x2,
                dc_12_prob=dc_12,
                dc_pick=new_dc_pick,
                dc_confidence=new_dc_conf
            )
            
        # Select best score that matches MS Pick
        # NEW LOGIC: We trust XGBoost/Poisson top score over generic MS Pick if MS Confidence is low.
        # Otherwise, we filter the grid to match the MS pick.
        ms_pick = reconciled_result.ms_pick
        
        def _score_matches_ms(score_str, pick):
            h, a = map(int, score_str.split("-"))
            if pick == "1": return h > a
            if pick == "2": return h < a
            return h == a
            
        matching_scores = [(s, p) for s, p in sorted_scores if _score_matches_ms(s, ms_pick)]
        
        # Primary Prediction Strategy:
        # If MS pick is highly confident, enforce it.
        # But if the absolute best score in the grid contradicts it and has a high probability (e.g. >10%), trust the score model directly.
        top_overall_score, top_overall_prob = sorted_scores[0]
        
        if matching_scores and not (top_overall_prob > 12.0 and not _score_matches_ms(top_overall_score, ms_pick)):
            predicted_ft = matching_scores[0][0]
        else:
            predicted_ft = top_overall_score
            
        # If we didn't calculate HT via ML (exception case), do it now
        if predicted_ht is None:
            ft_to_ht = self.config.get("half_time.ft_to_ht_ratio", 0.42)
            ht_h = round(lambda_home * ft_to_ht)
            ht_a = round(lambda_away * ft_to_ht)
            predicted_ht = f"{ht_h}-{ht_a}"
        
        # --- CONSISTENCY CHECK ---
        # Ensure HT score <= FT score
        try:
            ft_h, ft_a = map(int, predicted_ft.split("-"))
            ht_h, ht_a = map(int, predicted_ht.split("-"))
            
            # Clamp HT values
            ht_h = min(ht_h, ft_h)
            ht_a = min(ht_a, ft_a)
            
            predicted_ht = f"{ht_h}-{ht_a}"
        except ValueError:
            pass # Malformed score string, ignore correction

        ft_scores = [{"score": s, "prob": p} for s, p in sorted_scores[:5]]
        
        return ScorePrediction(
            predicted_ft_score=predicted_ft,
            predicted_ht_score=predicted_ht,
            ft_scores_top5=ft_scores,
            reconciled_ms=reconciled_result
        )
