"""Reproduce Pluraterra research 0.2 essential-access bounds.

Requires Python 3, numpy, scipy. Download essential-lab.json into the same
folder, then run: python reproduce-essentials.py

The JSON contains the exact frozen national marginals, scenario weights,
state order, results, and feasible endpoint distributions. Bounds are
conditional on those estimates and the stated model, not confidence intervals.
"""
import json
from pathlib import Path
import numpy as np
from scipy.optimize import linprog

data=json.loads(Path(__file__).with_name('essential-lab.json').read_text())
states=np.array(data['state_order'],dtype=float)
a_eq=np.vstack([np.ones(len(states)),states.T])
solves=0
for row in data['results']:
 b_eq=np.r_[1.,row['marginals']]
 for scenario,weights in data['scenarios'].items():
  utility=np.sqrt(states@np.array(weights))
  low=linprog(utility,A_eq=a_eq,b_eq=b_eq,bounds=(0,None),method='highs')
  high=linprog(-utility,A_eq=a_eq,b_eq=b_eq,bounds=(0,None),method='highs')
  assert low.success and high.success,(row['country'],scenario)
  assert np.max(np.abs(a_eq@low.x-b_eq))<1e-7
  assert np.max(np.abs(a_eq@high.x-b_eq))<1e-7
  actual=[100*low.fun**2,100*high.fun**2]
  expected=[row['scenarios'][scenario]['lower'],row['scenarios'][scenario]['upper']]
  assert np.max(np.abs(np.array(actual)-expected))<1e-5,(row['country'],scenario)
  solves+=2
print(f'Reproduced {solves} optimization endpoints across {len(data["results"])} economies.')