import pandas as pd
import numpy as np

class FeatureEngineer:
    def __init__(self, rankings_df, form_df, teams_df, manual_features_df=None):
        self.rankings = rankings_df
        self.form = form_df
        self.teams = teams_df
        self.manual_features = manual_features_df if manual_features_df is not None else pd.DataFrame()
        
        self.team_momentum = self._calculate_momentum()

    def _calculate_momentum(self):
        """
        Calculate a form momentum score for each team based on their last 10 matches.
        A win against a strong opponent (high Elo) gives more momentum than against a weak one.
        """
        momentum = {}
        for team_id, group in self.form.groupby('team_id'):
            # Take up to 10 most recent matches
            recent = group.head(10).copy()
            if len(recent) == 0:
                momentum[team_id] = 0
                continue
                
            # Points: 3 for win, 1 for draw, 0 for loss
            recent['points'] = np.where(recent['goals_for'] > recent['goals_against'], 3,
                               np.where(recent['goals_for'] == recent['goals_against'], 1, 0))
            
            # Weight points by opponent Elo (normalized roughly around 1500)
            # E.g., beating an 1800 Elo team gives weight 1.2, a 1200 Elo team gives 0.8
            recent['weight'] = recent['opponent_elo'] / 1500.0
            
            # Weighted average points
            score = (recent['points'] * recent['weight']).sum() / recent['weight'].sum()
            momentum[team_id] = score
            
        return pd.Series(momentum, name='momentum')

    def engineer_match_features(self, match_id, home_team_id, away_team_id):
        """
        Given match_id and home/away team IDs, return a feature dictionary.
        """
        # Elo rating
        if home_team_id not in self.rankings.index:
            raise ValueError(f"Missing ranking for home team: {home_team_id}")
        if away_team_id not in self.rankings.index:
            raise ValueError(f"Missing ranking for away team: {away_team_id}")
            
        home_rating = self.rankings.loc[home_team_id, 'rating']
        away_rating = self.rankings.loc[away_team_id, 'rating']
        
        if pd.isna(home_rating) or pd.isna(away_rating):
            raise ValueError(f"NaN rating for {home_team_id} or {away_team_id}")
        
        # Momentum
        home_mom = self.team_momentum.get(home_team_id, 1.0)
        away_mom = self.team_momentum.get(away_team_id, 1.0)
        
        # Host advantage
        home_host = 1 if (home_team_id in self.teams.index and self.teams.loc[home_team_id, 'is_host']) else 0
        away_host = 1 if (away_team_id in self.teams.index and self.teams.loc[away_team_id, 'is_host']) else 0
        
        # In World Cup, if home team is host, they have home advantage.
        # Often neither is host, meaning neutral venue.
        home_advantage = 1 if home_host else 0
        
        # Manual Features Extraction
        injury_impact_home = 0.0
        injury_impact_away = 0.0
        lineup_strength_home = 1.0
        lineup_strength_away = 1.0
        odds_implied_home_prob = 0.0
        odds_implied_away_prob = 0.0
        
        if match_id in self.manual_features.index:
            m_feat = self.manual_features.loc[match_id]
            injury_impact_home = m_feat.get('injury_impact_home') or 0.0
            injury_impact_away = m_feat.get('injury_impact_away') or 0.0
            lineup_strength_home = m_feat.get('lineup_strength_home') or 1.0
            lineup_strength_away = m_feat.get('lineup_strength_away') or 1.0
            
            odds_h = m_feat.get('odds_1x2_home')
            odds_d = m_feat.get('odds_1x2_draw')
            odds_a = m_feat.get('odds_1x2_away')
            
            # Simple implied probability from odds (1/odds)
            if pd.notna(odds_h) and pd.notna(odds_d) and pd.notna(odds_a):
                margin = (1/odds_h) + (1/odds_d) + (1/odds_a)
                odds_implied_home_prob = (1/odds_h) / margin
                odds_implied_away_prob = (1/odds_a) / margin

        return {
            'elo_diff': home_rating - away_rating,
            'momentum_diff': home_mom - away_mom,
            'home_advantage': home_advantage,
            'injury_impact_home': float(injury_impact_home),
            'injury_impact_away': float(injury_impact_away),
            'lineup_strength_home': float(lineup_strength_home),
            'lineup_strength_away': float(lineup_strength_away),
            'odds_implied_home_prob': float(odds_implied_home_prob),
            'odds_implied_away_prob': float(odds_implied_away_prob)
        }

    def build_dataset(self, matches_df):
        """
        Build a training dataset from a matches dataframe.
        """
        features = []
        labels = []
        
        for _, row in matches_df.iterrows():
            if pd.isna(row['home_team_id']) or pd.isna(row['away_team_id']):
                continue
            
            try:
                # Extract features
                f = self.engineer_match_features(row['id'], row['home_team_id'], row['away_team_id'])
                # Track the original index to match back easily
                f['match_id'] = row['id']
                features.append(f)
                
                # Determine label (0: Away Win, 1: Draw, 2: Home Win)
                if not pd.isna(row['home_score_90']) and not pd.isna(row['away_score_90']):
                    if row['home_score_90'] > row['away_score_90']:
                        labels.append(2)
                    elif row['home_score_90'] == row['away_score_90']:
                        labels.append(1)
                    else:
                        labels.append(0)
                else:
                    labels.append(None) # Unplayed
            except ValueError as e:
                print(f"Skipping match {row['id']}: {e}")
                
        return pd.DataFrame(features), pd.Series(labels)