import pandas as pd
import numpy as np
from pulp import *
from typing import List, Dict, Tuple

class DFSOptimizer:
    def __init__(self):
        self.model = None
        self.player_vars = {}
        self.constraints = []
        
    def setup_model(self, players_df: pd.DataFrame, salary_cap: float = 50000):
        """Initialize the optimization model with basic constraints"""
        self.model = LpProblem("NBA_DFS_Optimizer", LpMaximize)
        
        # Create binary variables for each player
        for idx, player in players_df.iterrows():
            var_name = f"player_{idx}"
            self.player_vars[var_name] = LpVariable(var_name, 0, 1, LpBinary)
        
        # Objective: Maximize projected points
        self.model += lpSum([self.player_vars[f"player_{idx}"] * row["Projected Points"]
                            for idx, row in players_df.iterrows()])
        
        # Constraint: Salary cap
        self.model += lpSum([self.player_vars[f"player_{idx}"] * row["Credits"]
                            for idx, row in players_df.iterrows()]) <= salary_cap
        
        # Constraint: Total players = 8
        self.model += lpSum(self.player_vars.values()) == 8
        
        # Position constraints
        self._add_position_constraints(players_df)
    
    def _add_position_constraints(self, players_df: pd.DataFrame):
        """Add NBA position constraints"""
        positions = {"PG": (1, 4), "SG": (1, 4), "SF": (1, 4), 
                    "PF": (1, 4), "C": (1, 4)}
        
        for pos, (min_count, max_count) in positions.items():
            pos_vars = [self.player_vars[f"player_{idx}"]
                       for idx, row in players_df.iterrows()
                       if row["Position"] == pos]
            if pos_vars:
                self.model += lpSum(pos_vars) >= min_count
                self.model += lpSum(pos_vars) <= max_count
    
    def add_team_stack(self, players_df: pd.DataFrame, min_players: int = 2, 
                       max_players: int = 5):
        """Add team stacking constraints"""
        teams = players_df["Team Name"].unique()
        for team in teams:
            team_vars = [self.player_vars[f"player_{idx}"]
                        for idx, row in players_df.iterrows()
                        if row["Team Name"] == team]
            if team_vars:
                self.model += lpSum(team_vars) >= min_players
                self.model += lpSum(team_vars) <= max_players
    
    def add_exposure_constraints(self, players_df: pd.DataFrame, 
                                max_exposure: float = 0.6):
        """Add player exposure constraints"""
        for player in players_df["Player"].unique():
            player_vars = [self.player_vars[f"player_{idx}"]
                          for idx, row in players_df.iterrows()
                          if row["Player"] == player]
            if player_vars:
                self.model += lpSum(player_vars) <= max_exposure
    
    def optimize(self) -> Tuple[List[str], float, float]:
        """Solve the optimization model and return the optimal lineup"""
        self.model.solve()
        
        if self.model.status != 1:
            raise Exception("Could not find optimal solution")
            
        selected_players = []
        total_points = 0
        total_salary = 0
        
        for var in self.model.variables():
            if var.varValue > 0:
                player_idx = int(var.name.split('_')[1])
                selected_players.append({
                    "Player": players_df.iloc[player_idx]["Player"],
                    "Position": players_df.iloc[player_idx]["Position"],
                    "Team": players_df.iloc[player_idx]["Team Name"],
                    "Salary": players_df.iloc[player_idx]["Credits"],
                    "Projected": players_df.iloc[player_idx]["Projected Points"]
                })
                total_points += players_df.iloc[player_idx]["Projected Points"]
                total_salary += players_df.iloc[player_idx]["Credits"]
        
        return selected_players, total_points, total_salary
    
    def generate_lineups(self, players_df: pd.DataFrame, num_lineups: int = 10,
                         min_salary: float = 48500, max_salary: float = 50000,
                         min_stack: int = 2, max_stack: int = 5,
                         max_exposure: float = 0.6) -> List[Dict]:
        """Generate multiple unique lineups"""
        lineups = []
        for _ in range(num_lineups):
            try:
                self.setup_model(players_df, max_salary)
                self.add_team_stack(players_df, min_stack, max_stack)
                self.add_exposure_constraints(players_df, max_exposure)
                
                # Add minimum salary constraint
                self.model += lpSum([self.player_vars[f"player_{idx}"] * row["Credits"]
                                    for idx, row in players_df.iterrows()]) >= min_salary
                
                # Add uniqueness constraint for previous lineups
                for prev_lineup in lineups:
                    prev_players = [p["Player"] for p in prev_lineup["players"]]
                    overlap_vars = [self.player_vars[f"player_{idx}"]
                                  for idx, row in players_df.iterrows()
                                  if row["Player"] in prev_players]
                    if overlap_vars:
                        self.model += lpSum(overlap_vars) <= 6  # Allow max 6 players overlap
                
                players, points, salary = self.optimize()
                lineups.append({
                    "players": players,
                    "total_points": points,
                    "total_salary": salary
                })
            except Exception as e:
                print(f"Failed to generate lineup: {str(e)}")
                continue
        
        return lineups