# ...existing code...

import numpy as np
import matplotlib.pyplot as plt

def voter_happiness(P, Pb, EFF, IMPmax, CTA, EXPEND, EXPENDnom):
    #Calculate the total priority and the cost per point for each policy
    Ptot = np.sum(P) + Pb #Total priority (sum of policy priorities and budget priority)
    CPP = CTA/IMPmax #Cost per point for each policy
    # print('CPP = ', CPP)
    EXPENDtot = EXPENDb = np.sum(EXPEND) #Total expenditure on policies
    #EXPENDtot = EXPENDb = 0.0
    IWM = EXPEND/CPP #Improvement with money -- how much improvement is achieved solely through expenditure
    # print('IWM = ', IWM)
    #The voter is most happy if no money is spent on any policy, thus giving IMPmax for the budget
    # beyond that happiness decreases linearly with expenditure
    IMPb = IMPmax - (IMPmax/EXPENDnom)*EXPENDtot #Improvement from budget
    IMP = IWM #initial IMP
    tolerance = 1e-6
    max_iterations = 1000
    for _ in range(max_iterations):
        new_IMP = IWM + np.dot(EFF, IMP)
        new_IMP = np.clip(new_IMP, None, IMPmax)  # Ensure IMP does not exceed 100
        if np.allclose(new_IMP, IMP, atol=tolerance):
            break
        IMP = new_IMP
    # print('Converged IMP = ', IMP)
    VOTHAP = IMP*P/Ptot #Voter's happiness from each policy
    # print('VOTHAP = ', VOTHAP)
    # print('IMPb = ', IMPb)
    VOTHAPb = IMPb*Pb/Ptot #Voter's happiness from budget
    # print('VOTHAPb = ', VOTHAPb)
    VOTHAPtot = np.sum(VOTHAP) + VOTHAPb
    # print('VOTHAPtot = ', VOTHAPtot)
    return VOTHAPtot

def find_pareto_frontier(costs, values):
    sorted_indices = np.argsort(costs)
    pareto_costs = []
    pareto_values = []
    max_value = -np.inf
    for idx in sorted_indices:
        cost = costs[idx]
        value = values[idx]
        if value > max_value:
            pareto_costs.append(cost)
            pareto_values.append(value)
            max_value = value
    return pareto_costs, pareto_values

#Manual USER INPUT
IMPmax = 100.0 #maximum improvement possible for each policy
P = np.array([10.0, 8.0, 10.0]) #Priority of each policy
Pb = 10.0 #Priority of budget
#Effects matrix -- how much each policy affects the others
#Effect of Column on Row -- EFF[Row][Col]
#Let's say the policies are Health, Education, and Infrastructure
#EFF[2][1] = 0.05 means that Education (Col 1) has a 5% effect on Infrastructure (Row 2)
EFF = np.array([[0, 0.10, 0.05],
                [0.10, 0, 0.1],
                [0.0, 0.05, 0.0]])

# EFF = np.array([[0, -0.10, 0.05],
#                 [0.10, 0, -0.1],
#                 [0.0, 0.05, 0.0]])

# print('EFF[0][1] =  ', EFF[0][1])
# print('EFF[1][0] =  ', EFF[1][0])
# print('EFF[2][1] =  ', EFF[2][1])
CTA = np.array([500.0, 400.0, 300.0]) #Cost to achieve max improvement for each policy
EXPEND  = np.array([297.96, 112.23, 266.61]) #Actual expenditure on each policy
EXPENDnom = 500.0 #Nominal expenditure (the expenditure at which voters are neutral)
#EXPEND = CTA
# print('EFF = ', EFF)

VOTHAPtot = voter_happiness(P, Pb, EFF, IMPmax, CTA, EXPEND, EXPENDnom)
print('VOTHAPtot = ', VOTHAPtot)

#Design of Experiments
# Vary the EXPEND values randomly and see how VOTHAPtot changes
num_points = 1000
EXPEND_ranges = [np.random.uniform(0, cta, num_points) for cta in CTA]

total_expenditures = []
vothap_totals = []

# Open file to write the table
with open('vothap_results.txt', 'w') as file:
    # Write headers
    file.write("EXPEND[0]\tEXPEND[1]\tEXPEND[2]\tTotalEXPEND\tVOTHAPtot\n")
    print("EXPEND[0]\tEXPEND[1]\tEXPEND[2]\tTotalEXPEND\tVOTHAPtot")
    
    for i in range(num_points):
        exp1 = EXPEND_ranges[0][i]
        exp2 = EXPEND_ranges[1][i]
        exp3 = EXPEND_ranges[2][i]
        EXPEND_test = np.array([exp1, exp2, exp3])
        total_expend = np.sum(EXPEND_test)
        VOTHAPtot = voter_happiness(P, Pb, EFF, IMPmax, CTA, EXPEND_test, EXPENDnom)
        result_line = f'{exp1}\t{exp2}\t{exp3}\t{total_expend}\t{VOTHAPtot}\n'
        file.write(result_line)
        print(result_line.strip())
        
        total_expenditures.append(total_expend)
        vothap_totals.append(VOTHAPtot)

# Find Pareto frontier
pareto_costs, pareto_values = find_pareto_frontier(total_expenditures, vothap_totals)

# Write Pareto frontier to a file
with open('pareto_frontier.txt', 'w') as pareto_file:
    pareto_file.write("TotalEXPEND\tVOTHAPtot\n")
    for cost, value in zip(pareto_costs, pareto_values):
        pareto_file.write(f'{cost}\t{value}\n')

# Plotting the results
plt.scatter(total_expenditures, vothap_totals, label='Total Results')
plt.scatter(pareto_costs, pareto_values, color='red', label='Pareto Frontier')
plt.xlabel('Total Expenditure')
plt.ylabel('Total Voter Happiness')
plt.title('Total Expenditure vs Total Voter Happiness')
plt.legend()
plt.grid(True)
plt.show()

#1. Vary the EXPEND values and see how VOTHAPtot changes