import streamlit as st
import pandas as pd
import numpy as np
import optuna

st.title('Project Budget Optimization')

# Define the costs for each project
costos = {
    'A': [120, 150, 180, 210, 240],
    'B': [180, 225, 270, 315, 360],
    'C': [200, 250, 300, 350, 400],
    'D': [240, 300, 360, 420, 480]
}

costs_df = pd.DataFrame(costos)

st.write("Costos de los proyectos a través de 5 años:")
st.write(costs_df)

# Get the discount rate and budget from the user
discount_rate = st.number_input("Ingrese la tasa de descuento (en formato decimal):", value=0.1)
budget = st.number_input("Ingrese el presupuesto total disponible:", value=1000)

# Calculate NPV for each project
def calculate_npv(costs, discount_rate):
    npv = sum(cost / ((1 + discount_rate) ** t) for t, cost in enumerate(costs, start=1))
    return npv

npvs = {project: calculate_npv(costs, discount_rate) for project, costs in costos.items()}
npvs_df = pd.DataFrame(list(npvs.items()), columns=['Project', 'NPV'])

st.write("Valor Presente Neto (VPN) de cada proyecto:")
st.write(npvs_df)

# Optimization
def objective(trial):
    # Suggest weights for each project (0 or 1 to include/exclude project)
    w_A = trial.suggest_int('w_A', 0, 1)
    w_B = trial.suggest_int('w_B', 0, 1)
    w_C = trial.suggest_int('w_C', 0, 1)
    w_D = trial.suggest_int('w_D', 0, 1)
    
    weights = np.array([w_A, w_B, w_C, w_D])
    
    # Calculate total cost and NPV
    total_cost = np.sum(weights * costs_df.sum(axis=0))
    total_npv = np.sum(weights * npvs_df['NPV'])
    
    # Penalize if total cost exceeds the budget
    if total_cost > budget:
        return -1e10
    
    return total_npv

study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=100, show_progress_bar=True)

# Extract best parameters
best_params = study.best_params
selected_projects = [project for project, selected in best_params.items() if selected]

st.write("Los mejores proyectos seleccionados son:", selected_projects)

# Calculate the total cost and NPV of the selected projects
total_cost = np.sum([costos[project] for project in selected_projects], axis=0).sum()
total_npv = np.sum([npvs[project] for project in selected_projects])

st.write(f"El costo total de los proyectos seleccionados es: :green[{total_cost}]")
st.write(f"El VPN total de los proyectos seleccionados es: :green[{total_npv}]")