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
from pyomo.environ import *

def get_output(df, df1, df2):
    df.fillna(0, inplace=True)
    df1.fillna(0, inplace=True)
    df2.fillna(0, inplace=True)
    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]]

    # 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)
    model.y = Var(flat_J, model.K, domain=Binary)
    model.s = Var(flat_J, domain=NonNegativeReals)

    # Objective function
    def objective_rule(model):
        return sum(model.p[i] * model.s[i, j] for i in model.I for j in model.J[i])
    model.objective = Objective(rule=objective_rule, sense=minimize)

    # 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)

    # Constraint to ensure each task is assigned to exactly one resource
    def single_resource_constraint(model, i, j):
        return sum(model.y[i, j, k] for k in model.K) == 1
    model.single_resource_constraint = Constraint(flat_J, rule=single_resource_constraint)

    # Linking x and y
    def linking_constraint(model, i, j, k, month):
        return model.x[i, j, k, month] <= 1000 * model.y[i, j, k]
    model.linking_constraint = Constraint(flat_J, model.K, model.M, rule=linking_constraint)

    # Ensure glissement plus capacité allouée égale à planifiée
    def glissement_constraint(model, i, j):
        return model.s[i, j] >= sum(model.H[i, j, m] for m in model.M) - sum(model.x[i, j, k, m] * model.A[i, j, k] for k in model.K for m in model.M)
    model.glissement_constraint = Constraint(flat_J, rule=glissement_constraint)

    # Ensure glissement is non-negative
    def non_negative_glissement_constraint(model, i, j):
        return model.s[i, j] >= 0
    model.non_negative_glissement_constraint = Constraint(flat_J, rule=non_negative_glissement_constraint)

    # Ensure x is less than or equal to H
    def x_less_than_H_constraint(model, i, j, k, m):
        return model.x[i, j, k, m] <= model.H[i, j, m]
    model.x_less_than_H_constraint = Constraint(flat_J, model.K, model.M, rule=x_less_than_H_constraint)

    # Solver
    solver = SolverFactory('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)
    df_finall = output_df.loc[output_df[Months].sum(axis=1) > 0]
    return df_finall



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()