from gekko import GEKKO
import numpy as np
import matplotlib.pyplot as plt

# Initialize the Gekko model
m = GEKKO(remote=True)

# Parameters
failure_rate = m.Param(value=0.01)      # Failure rate per day
cost_maintenance = m.Param(value=1000)  # Cost per maintenance
cost_failure = m.Param(value=5000)      # Cost per failure

# Variables
interval = m.FV(value=90, lb=30, ub=365)  # Maintenance interval (days)
interval.STATUS = 1  # Allow optimization

# Maintenance cost function: Periodic maintenance cost
maintenance_cost = m.Intermediate(cost_maintenance * (365 / interval))

# Failure cost function: Failures that happen if no maintenance
failures = m.Intermediate(failure_rate * (interval / 2))  # Average failures
failure_cost = m.Intermediate(cost_failure * failures)

# Total cost: Sum of maintenance cost and failure cost
total_cost = m.Var()
m.Equation(total_cost==maintenance_cost + failure_cost)

# Objective: Minimize the total cost
m.Minimize(total_cost)

# Set solver options
m.options.IMODE = 3  # Steady-state optimization
m.options.SOLVER = 1 # APOPT Solver

# Solve the optimization problem
m.solve(disp=True)

# Print optimized maintenance interval
opt = interval.value[0]
print(f"Optimized maintenance interval: {opt:.2f} days")

# Create an array of intervals to calculate total costs
intervals = np.linspace(30, 365, 50)
total_costs = []
maintenance_costs = []
failure_costs = []

interval.STATUS = 0
for i in intervals:
    interval.value = i
    m.solve(disp=False)  # Solve for each interval value
    total_costs.append(total_cost.value[0])
    maintenance_costs.append(maintenance_cost.value[0])
    failure_costs.append(failure_cost.value[0])

# Plot total cost vs. interval
plt.figure(figsize=(7, 4))
plt.plot(intervals, total_costs, label='Total Cost')
plt.plot(intervals, maintenance_costs, label='Maintenance Cost')
plt.plot(intervals, failure_costs, label='Failure Cost')
plt.axvline(opt, color='red', linestyle='--',
            label=f'Optimized Interval: {opt:.2f} days')
plt.xlabel('Maintenance Interval (days)')
plt.ylabel('Total Cost ($)')
plt.legend(); plt.grid(True)
plt.tight_layout()
plt.savefig('total_cost_vs_interval.png', dpi=300)
plt.show()