import numpy as np
import pandas as pd
from mip import Model, xsum, maximize, BINARY


def single_game_df(csv):
    """
    csv : csv as generated by DraftKings
    """
    df = pd.read_csv(csv)
    df = df[df['Roster Position'] == 'FLEX']
    df = df[df['Salary'] > 200]
    df = df.reset_index(drop=True)
    return df


def single_game(df, weight_col='AvgPointsPerGame', num_lineups=1):
    """
    df : dataframe from csv as generated by DraftKings
    weight_col : column name of reward func
    num_lineups : number of lineups to generate
    """
    N = len(df)
    # reward
    r = df[weight_col]
    # salary
    s = df['Salary']

    salary_cap = 50000
    I = range(N)

    model = Model("daily_fantasy")
    model.verbose = 0

    # players
    p = [model.add_var(var_type=BINARY) for i in I]
    # captain
    c = [model.add_var(var_type=BINARY) for i in I]

    model.objective = maximize(xsum(r[i] * (p[i] + 1.5 * c[i]) for i in I))

    model += xsum(s[i] * (p[i] + 1.5 * c[i]) for i in I) <= salary_cap
    for i in I:
        model += p[i] + c[i] <= 1
    model += xsum(p[i] for i in I) == 5
    model += xsum(c[i] for i in I) == 1

    lineups = []
    lineup = []
    for l in range(num_lineups):
        # Ensure that repeat playlists are not picked
        if l != 0 and len(lineup) != 0:
            a = np.zeros(N)
            a[lineup[0]] = 2
            b = np.zeros(N)
            b[lineup[1:]] = 1
            model += xsum(a[i] * c[i] + b[i] * p[i] for i in I) <= 2 + (5 * 1) - 1

        model.optimize()
        lineup = [i for i in I if c[i].x >= 0.99] + [i for i in I if p[i].x >= 0.99]
        lineups.append(lineup)

    return lineups


def print_lineups(lineups, df, weight_col='AvgPointsPerGame'):
    for lineup in lineups:
        cost = 1.5 * df['Salary'][lineup[0]] + sum(df['Salary'][i] for i in lineup[1:])
        pts = 1.5 * df[weight_col][lineup[0]] + sum(df[weight_col][i] for i in lineup[1:])

        line = f"[{df['Name'][lineup[0]]}]"
        for p in lineup[1:]:
            line += f" - {df['Name'][p]}"
        print(f"Cost: {int(cost)}; Proj: {pts:.1f}")
        print(line)
