import streamlit as st
import pandas as pd
import io
from pyomo.environ import ConcreteModel, Var, Objective, Constraint, SolverFactory, NonNegativeReals, RangeSet, Param, minimize, value, Reals,Set

def get_output(df,df1,df2):
  n = df['ID projet'].nunique()
  task = df.groupby('ID projet').count()['Nom projet']
  project =df.groupby('ID projet').count().index

  J_sizes = {i:task[i-1] for i in range(1,n+1)}

  Months = ['Janvier', 'Février', 'Mars', 'Avril', 'Mai', 'Juin', 'Juillet', 'Août', 'Septembre', 'Octobre', 'Novembre', 'Décembre']
  months = 12  # Number of months in set M
  H_data = {(i, j, month): df.loc[df['ID projet']==project[i-1]].loc[df.loc[df['ID projet']==project[i-1]].index[j-1],Months[month-1]] for i in range(1, n + 1) for j in range(1, J_sizes[i] + 1) for month in range(1, months + 1)}

  df1.fillna(0, inplace=True)

  h = df1['Ressource'].nunique()
  A_data = {(i, j, k): int(df.loc[df['ID projet']==project[i-1]].loc[df.loc[df['ID projet']==project[i-1]].index[j-1], 'Equipe']  == df1.loc[df1.index[k-1],'Equipe']) for i in range(1, n + 1) for j in range(1, J_sizes[i] + 1) for k in range(1, h + 1)}
  per= [0.08,0.08,0.09,0.09,0.08,0.09,0.07,0.07,0.09,0.09,0.09,0.08]
  C_data = {(k, month): df1.loc[df1.index[k-1],'Capacité']*per[month-1] for k in range(1, h + 1) for month in range(1, months + 1)}

  p_data = {i:df2.loc[df2.index[i-1],'Pond']  for i in range(1, n + 1)}

    # Define model
  model = ConcreteModel()

  # Sets
  model.I = RangeSet(1, n)
  model.M = RangeSet(1, months)
  model.K = RangeSet(1, h)
  model.J = Set(model.I, initialize=lambda model, i: RangeSet(1, J_sizes[i]))

  # Flatten J for use in parameter definition
  flat_J = [(i, j) for i in model.I for j in model.J[i]]
  flat_J_pairs = [(i, j, l) for i in model.I for j in model.J[i] for l in model.J[i] if j != l]

  # Parameters
  model.H = Param(flat_J, model.M, initialize=H_data)
  model.A = Param(flat_J, model.K, initialize=A_data)
  model.C = Param(model.K, model.M, initialize=C_data)
  model.p = Param(model.I, initialize=p_data)

  # Variables
  model.x = Var(flat_J, model.K, model.M, domain=NonNegativeReals)

  # Auxiliary variables for max(0, ...)
  model.max_0_terms = Var(flat_J, model.M, domain=NonNegativeReals)
  model.max_0_terms_2 = Var(flat_J_pairs, model.M, domain=NonNegativeReals)

  # Objective function
  def objective_rule(model):
      return sum(model.p[i] * (
          sum(model.max_0_terms[i, j, month] for j in model.J[i]) +
          sum(model.max_0_terms_2[i, j, l, month] for j in model.J[i] for l in model.J[i] if l != j)
      ) for i in model.I for month in model.M)
  model.objective = Objective(rule=objective_rule, sense=minimize)

  # Constraints to handle max(0, ...)
  def max_0_term_constraint_1(model, i, j, month):
      return model.max_0_terms[i, j, month] >= model.H[i, j, month] - sum(model.A[i, j, k] * model.x[i, j, k, month] for k in model.K)
  model.max_0_term_constraint_1 = Constraint(flat_J, model.M, rule=max_0_term_constraint_1)

  def max_0_term_constraint_2(model, i, j, month):
      return model.max_0_terms[i, j, month] >= 0
  model.max_0_term_constraint_2 = Constraint(flat_J, model.M, rule=max_0_term_constraint_2)

  def max_0_term_2_constraint_1(model, i, j, l, month):
      return model.max_0_terms_2[i, j, l, month] >= model.H[i, l, month] - sum(model.A[i, l, k] * model.x[i, l, k, month] for k in model.K)
  model.max_0_term_2_constraint_1 = Constraint(flat_J_pairs, model.M, rule=max_0_term_2_constraint_1)

  def max_0_term_2_constraint_2(model, i, j, l, month):
      return model.max_0_terms_2[i, j, l, month] >= 0
  model.max_0_term_2_constraint_2 = Constraint(flat_J_pairs, model.M, rule=max_0_term_2_constraint_2)

  # Capacity constraint
  def capacity_constraint(model, k, month):
      return sum(model.x[i, j, k, month] for (i, j) in flat_J) <= model.C[k, month]
  model.capacity_constraint = Constraint(model.K, model.M, rule=capacity_constraint)

  # Solver
  solver = SolverFactory('glpk')  # Example using GLPK
  result = solver.solve(model, tee=True)

  Months= ['Janvier', 'Février', 'Mars', 'Avril', 'Mai', 'Juin', 'Juillet', 'Août', 'Septembre', 'Octobre', 'Novembre', 'Décembre']
  results = []
  for (i, j) in flat_J:
      for k in model.K:
        result = {}
        result['i']=project[i-1]
        result['j']= j
        result['k']=df1.loc[df1.index[k-1],'Ressource']
        for month in model.M:
                
                result[Months[month-1]] = value(model.x[i, j, k, month])
        results.append(result)


  output_df = pd.DataFrame(results)
  return output_df



def main():
    st.title("XLSX Upload and Download")

    # File upload section
    uploaded_file = st.file_uploader("Choose an XLSX file to upload", type="xlsx")

    if uploaded_file is not None:
        # Load the uploaded file into a Pandas DataFrame
        df = pd.read_excel(uploaded_file, sheet_name='PMC1')
        df1 = pd.read_excel(uploaded_file, sheet_name ='Base de ressource1')
        df2 =  pd.read_excel(uploaded_file, sheet_name ='Priorisation')

        df_out = get_output(df,df1,df2)
        # Display the uploaded DataFrame
        st.write("Estimation")
        st.dataframe(df_out)

        # Download section
        excel_file = io.BytesIO()
        df_out.to_excel(excel_file, index=False)

        st.download_button(
            label="Download XLSX",
            data=excel_file.getvalue(),
            file_name="downloaded_file.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        )

if __name__ == "__main__":
    main()