import pandas as pd
import numpy as np
from pulp import LpProblem, LpVariable, LpMaximize, LpMinimize, LpStatus, lpSum, value, GLPK_CMD, PULP_CBC_CMD
import random

## Feed Mills
feed_mill = ["Mill1","Mill2","Mill3","Mill4"]
## Diet Types
diet_types = ["NUR1","NUR2","FIN1","FIN2","FIN3","LAC","GES"]
## baseline_costlb
# Baseline costs / lb by Diet Type
baseline_costs = {"NUR1": 1.60,"NUR2": 1.30,"FIN1": 1.10,"FIN2": 1.00,"FIN3": 0.85,"LAC": 1.35,"GES": 0.95}
# Difference in costs by mill
feedmill_costs = {"Mill1": -0.20,"Mill2": -0.01,"Mill3":0.15,"Mill4":0.10}
## Cost_lb by diet type and mill
costs_long = []
for diet in diet_types:
    for feed in feed_mill:
        cost = round(baseline_costs[diet] + feedmill_costs[feed] + random.uniform(-0.05, 0.05), 2)
        costs_long.append({"feed_mill": feed, "diet_type": diet, "cost_lb": cost})

costs_df = pd.DataFrame(costs_long)
print(costs_df)
## Number of Orders (each order is a farm site)
loads = 25
## Simulate 'n' Feed Orders 
orders = []
for i in range(loads):
    if random.random() < 0.75:
        feeds = [random.choice(diet_types)]  # Full order of a single feed type
    else:
        group = random.choice([["NUR1","NUR2"],["FIN1","FIN2","FIN3"],["LAC","GES"]])
        feeds = random.sample(group, random.randint(2, min(3, len(group)))) #split order
    orders.append({'order': i+1, 'diet_type': feeds})
## turn it into pandas dataframe and add feed lbs
order_df = (pd.DataFrame(orders)
    .explode('diet_type')
    #.rename(columns={'Feeds': 'diet_type'})
    .reset_index(drop=True)
    .assign(count = lambda x: x.groupby('order').transform('count'))
    .assign(feedlbs = lambda x: round(26000 / x['count'],0))
    .filter(['order','diet_type','feedlbs'])
)
print(order_df)
## randomly create mileage data
mileage_data = {'order': [],'feed_mill': [],'miles': []}
for i in range(loads):
    for mill in feed_mill:
        mileage_data['order'].append(i+1)
        mileage_data['feed_mill'].append(mill)
        mileage_data['miles'].append(random.randint(5, 100))  # Random miles between 5 and 100
mileage_df = pd.DataFrame(mileage_data)
print(mileage_df)

## build cost matrix  
cost_matrix = (order_df
    ## take every order to every mill
    .merge(pd.DataFrame(costs_df['feed_mill'].drop_duplicates()), how='cross')
    ## bring in miles from feed mill to grower
    .merge(mileage_df,on=['order','feed_mill'], how='left')
    ## bring diet cost for each mill
    .merge(costs_df, on=['feed_mill','diet_type'], how='left')
   ## get costs
    .assign(cost_miles = lambda x: (x['miles'] * 4.0))
    .assign(cost_feed = lambda x: (x['feedlbs'] * x['cost_lb']))
    .assign(cost = lambda x: (x['cost_feed'] + x['cost_miles']))
    .assign(cost_lb = lambda x: x['cost'] / x['feedlbs'])
    .filter(['order','diet_type','feed_mill','cost_lb'])
)

mill_max = (costs_df[['feed_mill']]
    .drop_duplicates()
    .assign(feedlbs = order_df['feedlbs'].sum()*0.40)
)

dietmill_max = (costs_df[['feed_mill','diet_type']]
    .drop_duplicates()
    .assign(feedlbs = order_df['feedlbs'].sum()*0.40)
    ## mill 4 is only sow farm feed
    .assign(feedlbs=lambda x: np.where((x['feed_mill'].isin(['Mill1','Mill2','Mill3'])) & (x['diet_type'] == 'LAC'), 0, x['feedlbs']))
    .assign(feedlbs=lambda x: np.where((x['feed_mill'].isin(['Mill4'])) & (x['diet_type'] == 'LAC'), order_df.query('diet_type == "LAC"').sum()['feedlbs'], x['feedlbs']))
    .assign(feedlbs=lambda x: np.where((x['feed_mill'].isin(['Mill1','Mill2','Mill3'])) & (x['diet_type'] == 'GES'), 0, x['feedlbs']))
    .assign(feedlbs=lambda x: np.where((x['feed_mill'].isin(['Mill4'])) & (x['diet_type'] == 'GES'), order_df.query('diet_type == "LAC"').sum()['feedlbs'], x['feedlbs']))
)



## set input
op_input_cost = cost_matrix.set_index(['order','diet_type','feed_mill']).to_dict()['cost_lb']
op_orders = set(zip(cost_matrix['order'],cost_matrix['diet_type'],cost_matrix['feed_mill']))
op_ordermill = set(zip(cost_matrix['order'],cost_matrix['feed_mill']))
op_orderdiet = set(zip(cost_matrix['order'],cost_matrix['diet_type']))       
op_dietmill = set(zip(dietmill_max['feed_mill'],dietmill_max['diet_type'])) 
op_demand = order_df.set_index(['order','diet_type']).to_dict()['feedlbs']
op_overall_supply = mill_max.set_index(['feed_mill']).to_dict()['feedlbs']
op_dietmill_supply = dietmill_max.set_index(['feed_mill','diet_type']).to_dict()['feedlbs']
op_mill = (cost_matrix.drop_duplicates('feed_mill')).feed_mill
op_order = (cost_matrix.drop_duplicates('order')).order

# 1. Initialise model
model = LpProblem("MinimizeFeedCosts", LpMinimize)
# 2. Define Decision Variables
deliveries =  LpVariable.dicts('order_mill', [(o, d, m) for o, d, m in op_orders], lowBound=0, cat='Continuous')
mill_assignment = LpVariable.dicts('assign_mill', [(o, d, m) for o, d, m in op_orders], cat='Binary')
# used to force orders to a single mill
mill_used_load = LpVariable.dicts('millused',[(o, m) for o, m in op_ordermill], cat='Binary')
total_mill = LpVariable.dicts('millcount',[(o) for o in op_order], cat='Continuous')
# 3. Define objective function: costs to ship feed for a load for a given feed mill
model += lpSum([op_input_cost[(o, d, m)] * deliveries[(o, d, m)] for o, d, m in op_orders])
# 4. Define constraints:
### Constraint 1: For each order diet demand met
for o, d in op_orderdiet:
    model += lpSum([deliveries[(o, d, m)] for m in op_mill]) == op_demand[(o, d)]
    # Add mill assignment constraint 
    model += lpSum([mill_assignment[(o, d, m)] for m in op_mill]) == 1
### Constraint 1.1: Deliveries for a diet type and order can only occur if corresponding mill is assigned 
for o, d in op_orderdiet:
    for m in op_mill:
        model += deliveries[(o, d, m)] <= op_demand[(o, d)] * mill_assignment[(o, d, m)]      
### Constraint 2: total feed from a mill has to be less than supply of a given mill
for m in op_mill:
    model += lpSum([deliveries[(o, d, m)] for o, d in op_orderdiet]) <= op_overall_supply[m]
### Constraint 3: total feed from a mill/diet has to be less than supply of a given mill/diet
for m, d in op_dietmill:
    tmp_order_comb = [(o) for o, d_inner, m_inner in op_orders if d == d_inner and m == m_inner]
    model += lpSum([deliveries[(o, d, m)] for o in tmp_order_comb]) <= op_dietmill_supply[m, d]
### Constraint 4: Ensure all diet types for an order come from the same mill
#Link processor load-suf assignments to processor assignments
for o, m in op_ordermill:
    tmp_diet_comb = [(d) for o_inner, d in op_orderdiet if o == o_inner]
    model +=(mill_used_load[(o, m)] * 5 >= (lpSum(mill_assignment[(o, d, m)] for d in tmp_diet_comb)))
    model +=(mill_used_load[(o, m)] <= lpSum(mill_assignment[(o, d, m)] for d in tmp_diet_comb))
    
# total processors for a load
for o in op_order:
    model += (total_mill[o] == lpSum([mill_used_load[(o, m)] for m in op_mill]))

## An order can only come from one mill
for o in op_order:
    model += (total_mill[o] == 1)

## optimize model
solver = PULP_CBC_CMD(threads=4, cuts=1, gapRel=0.01)
model.solve(solver)

variable_values = [(var.name, var.value()) for var in model.variables()]
solution = pd.DataFrame(variable_values, columns=["Variable", "Value"])
solution = solution[solution['Variable'].str.contains('assign_mill')]
solution = (solution
    .assign(order = (solution['Variable'].str.split(",").str[0]).str.replace("assign_mill_(","").astype(int),
            diet_type = (solution['Variable'].str.split(",_'").str[1]).str.replace("'",""),
            feed_mill = (solution['Variable'].str.split(",_").str[2]).str.replace("'","").str.replace(")",""))
    .filter(['order','diet_type','feed_mill','Value'])
    .query('Value == 1.0')
    .merge(order_df,on=['order','diet_type'],how='left')
    .sort_values(by=['order'])
    .filter(['order','diet_type','feed_mill','feedlbs'])
)

solution