pruncalc/roi.py

from __future__ import annotations

import dataclasses
import json
import typing

import cache

def main() -> None:
	recipes: list[Recipe] = cache.get('https://api.prunplanner.org/data/recipes/')
	buildings: dict[str, Building] = {m['building_ticker']: m for m in cache.get('https://api.prunplanner.org/data/buildings/')}
	materials: dict[str, Material] = {m['ticker']: m for m in cache.get('https://api.prunplanner.org/data/materials/')}
	raw_prices: list[RawPrice] = cache.get('https://refined-prun.github.io/refined-prices/all.json')
	
	# EXTREME DETAIL: We execute the "Open Source Heist" here. We pull the static JSON 
	# directly from PRUNplanner's GitHub repository. 
	hq_levels_raw = cache.get('https://raw.githubusercontent.com/prunplanner/prunplanner/main/frontend/src/features/hq_upgrade_calculator/hq_levels.json')
	
	for cx in ['AI1', 'CI1', 'IC1', 'NC1']:
		profits = calc_for_cx(cx, recipes, buildings, materials, raw_prices, hq_levels_raw)
		with open(f'www/roi_{cx.lower()}.json', 'w') as f:
			json.dump([dataclasses.asdict(p) for p in profits], f, indent='\t')

def calc_for_cx(cx: str, recipes: typing.Collection[Recipe], buildings: typing.Mapping[str, Building],
		materials: typing.Mapping[str, Material], raw_prices: typing.Collection[RawPrice], 
		hq_levels_raw: dict) -> typing.Sequence[Profit]:
	prices: dict[str, Price] = {
		p['MaterialTicker']: Price(p['VWAP7D'], p['AverageTraded7D'], p['VWAP30D'], p['Bid'], p['Ask']) for p in raw_prices # pyright: ignore[reportArgumentType]
		if p['ExchangeCode'] == cx
	}
	
	# EXTREME DETAIL: We pre-calculate the VWAP, Bid, and Ask prices for every HQ level.
	# This ensures the frontend doesn't have to do any heavy lifting, and the HQ CapEx
	# naturally inherits the user's active price-metric permutation from the dropdowns!
	hq_costs: dict[str, dict[str, float]] = {}
	for level_str, mats in hq_levels_raw.items():
		cost = {'vwap': 0.0, 'bid': 0.0, 'ask': 0.0}
		for mat in mats:
			if mat['ticker'] in prices:
				m_metrics = get_metrics(mat['amount'], prices[mat['ticker']])
				for k in cost: cost[k] += m_metrics[k]
		hq_costs[level_str] = cost

	habitation: typing.Mapping[Worker, str] = {
		'pioneers': 'HB1',
		'settlers': 'HB2',
		'technicians': 'HB3',
		'engineers': 'HB4',
		'scientists': 'HB5',
	}
	hab_area_cost: dict[Worker, float] = {}
	hab_capex: dict[Worker, dict[str, float]] = {}
	for worker, hab in habitation.items():
		hab_area_cost[worker] = buildings[hab]['area_cost'] / 100
		base_capex = building_construction_cost(buildings[hab], prices)
		hab_capex[worker] = {k: v / 100 for k, v in base_capex.items()}

	profits: list[Profit] = []
	for recipe in recipes:
		if profit := calc_profit(recipe, buildings, hab_area_cost, hab_capex, materials, prices, hq_costs):
			profits.append(profit)
	profits.sort()
	return profits

def get_metrics(amount: float, price: Price) -> dict[str, float]:
	v = price.vwap_7d or price.vwap_30d or 0.0
	b = price.bid if price.bid is not None else v
	a = price.ask if price.ask is not None else v
	return {'vwap': amount * v, 'bid': amount * b, 'ask': amount * a}

def calc_profit(recipe: Recipe, buildings: typing.Mapping[str, Building], hab_area_cost: typing.Mapping[Worker, float],
		hab_capex: typing.Mapping[Worker, dict[str, float]], materials: typing.Mapping[str, Material],
		prices: typing.Mapping[str, Price], hq_costs: dict[str, dict[str, float]]) -> Profit | None:
	if len(recipe['outputs']) == 0:
		return

	building = buildings[recipe['building_ticker']]
	area = building['area_cost'] + sum(hab_area_cost[worker] * building[worker] for worker in hab_area_cost)
	runs_per_day = 24 * 60 * 60 * 1000 / recipe['time_ms'] * 1.25 # assume CoGC
	if building['building_ticker'] in ('FRM', 'ORC'):
		runs_per_day *= 1.1212 # promitor's fertility

	outputs: list[MatPrice] = []
	revenue = {'vwap': 0.0, 'bid': 0.0, 'ask': 0.0}
	output_prices: dict[str, PriceNonNull] = {}
	for output in recipe['outputs']:
		price = prices[output['material_ticker']]
		if price.vwap_7d is None or price.average_traded_7d is None:
			return # skip recipes with thinly traded outputs
		output_prices[output['material_ticker']] = typing.cast(PriceNonNull, price)
		
		m = get_metrics(output['material_amount'] * runs_per_day, price)
		for k in revenue: revenue[k] += m[k]
		outputs.append(MatPrice(output['material_ticker'], output['material_amount'], price.vwap_7d, price.bid, price.ask))

	input_costs: list[MatPrice] = []
	opex = {'vwap': 0.0, 'bid': 0.0, 'ask': 0.0}
	for input in recipe['inputs']:
		price = prices[input['material_ticker']]
		if price.vwap_7d is None:
			return # skip recipes with thinly traded inputs
		
		m = get_metrics(input['material_amount'] * runs_per_day, price)
		for k in opex: opex[k] += m[k]
		input_costs.append(MatPrice(input['material_ticker'], input['material_amount'], price.vwap_7d, price.bid, price.ask))

	worker_consumable = building_daily_cost(building, prices)
	for k in opex: opex[k] += worker_consumable[k]

	capex = building_construction_cost(building, prices)
	for worker, hab_cost in hab_capex.items():
		workers = building[worker]
		if workers > 0:
			for k in capex: capex[k] += hab_cost[k] * workers

	lowest_liquidity = min(recipe['outputs'],
			key=lambda output: output['material_amount'] / output_prices[output['material_ticker']].average_traded_7d)
	output_per_day = lowest_liquidity['material_amount'] * runs_per_day
	
	average_traded_7d = output_prices[lowest_liquidity['material_ticker']].average_traded_7d
	
	output_per_base = output_per_day / (area / 500)
	market_capacity_base = average_traded_7d / output_per_base

	in_w = sum(materials[input['material_ticker']]['weight'] * input['material_amount'] for input in recipe['inputs'])
	in_v = sum(materials[input['material_ticker']]['volume'] * input['material_amount'] for input in recipe['inputs'])
	out_w = sum(materials[output['material_ticker']]['weight'] * output['material_amount'] for output in recipe['outputs'])
	out_v = sum(materials[output['material_ticker']]['volume'] * output['material_amount'] for output in recipe['outputs'])

	runs_per_base = runs_per_day / (area / 500)
	
	ships_needed_per_base = max(in_w / 3000, in_v / 1000, out_w / 3000, out_v / 1000) * runs_per_base
	ship_capex_per_base = ships_needed_per_base * 800_000

	bottlenecks = [
		(in_w, 't (I)'),
		(in_v, 'm³ (I)'),
		(out_w, 't (O)'),
		(out_v, 'm³ (O)')
	]
	max_logistics, logistics_bottleneck = max(bottlenecks, key=lambda x: x[0])
	logistics_per_base = max_logistics * runs_per_base
	
	return Profit(outputs, recipe['recipe_name'],
			expertise=building['expertise'],
			building=building['building_ticker'],
			area=area,
			capex=capex,
			opex=opex,
			revenue=revenue,
			input_costs=input_costs,
			runs_per_day=runs_per_day,
			logistics_per_base=logistics_per_base,
			logistics_bottleneck=logistics_bottleneck,
			output_per_day=output_per_day,
			average_traded_7d=average_traded_7d,
			market_capacity_base=market_capacity_base,
			ship_capex_per_base=ship_capex_per_base,
			hq_costs=hq_costs)

def building_construction_cost(building: Building, prices: typing.Mapping[str, Price]) -> dict[str, float]:
	cost = {'vwap': 0.0, 'bid': 0.0, 'ask': 0.0}
	for bc in building['costs']:
		m = get_metrics(bc['material_amount'], prices[bc['material_ticker']])
		for k in cost: cost[k] += m[k]
	# https://handbook.apex.prosperousuniverse.com/wiki/building-costs/#rocky-planets
	mcg = get_metrics(building['area_cost'] * 4, prices['MCG'])
	for k in cost: cost[k] += mcg[k]
	return cost

def building_daily_cost(building: Building, prices: typing.Mapping[str, Price]) -> dict[str, float]:
	consumption = {
		'pioneers': [('COF', 0.5), ('DW', 4), ('RAT', 4), ('OVE', 0.5), ('PWO', 0.2)],
		'settlers': [('DW', 5), ('RAT', 6), ('KOM', 1), ('EXO', 0.5), ('REP', 0.2), ('PT', 0.5)],
		'technicians': [('DW', 7.5), ('RAT', 7), ('ALE', 1), ('MED', 0.5), ('SC', 0.1), ('HMS', 0.5), ('SCN', 0.1)],
		'engineers': [('DW', 10), ('MED', 0.5), ('GIN', 1), ('FIM', 7), ('VG', 0.2), ('HSS', 0.2), ('PDA', 0.1)],
		'scientists': [('DW', 10), ('MED', 0.5), ('WIN', 1), ('MEA', 7), ('NST', 0.1), ('LC', 0.2), ('WS', 0.05)],
	}
	cost = {'vwap': 0.0, 'bid': 0.0, 'ask': 0.0}
	for worker, mats in consumption.items():
		workers = building[worker]
		for mat, per_100 in mats:
			m = get_metrics(workers * per_100 / 100, prices[mat])
			for k in cost: cost[k] += m[k]
	return cost

Worker = typing.Literal['pioneers', 'settlers', 'technicians', 'engineers', 'scientists']

class Recipe(typing.TypedDict):
	recipe_name: str
	building_ticker: str
	inputs: list[RecipeMat]
	outputs: list[RecipeMat]
	time_ms: int

class RecipeMat(typing.TypedDict):
	material_ticker: str
	material_amount: int

class Building(typing.TypedDict):
	building_ticker: str
	expertise: str
	area_cost: int
	costs: list[BuildingMat]
	pioneers: int
	settlers: int
	technicians: int
	engineers: int
	scientists: int

class BuildingMat(typing.TypedDict):
	material_ticker: str
	material_amount: int

class Material(typing.TypedDict):
	ticker: str
	weight: float
	volume: float

class RawPrice(typing.TypedDict):
	MaterialTicker: str
	ExchangeCode: str
	VWAP7D: float | None
	AverageTraded7D: float | None
	VWAP30D: float | None
	Bid: float | None
	Ask: float | None

@dataclasses.dataclass(eq=False, frozen=True, slots=True)
class Price:
	vwap_7d: float | None
	average_traded_7d: float | None
	vwap_30d: float | None
	bid: float | None
	ask: float | None

@dataclasses.dataclass(eq=False, frozen=True, slots=True)
class PriceNonNull:
	vwap_7d: float
	average_traded_7d: float

@dataclasses.dataclass(eq=False, frozen=True, slots=True)
class Profit:
	outputs: typing.Collection[MatPrice]
	recipe: str
	expertise: str
	building: str
	area: float
	capex: dict[str, float]
	opex: dict[str, float]
	revenue: dict[str, float]
	input_costs: typing.Collection[MatPrice]
	runs_per_day: float
	logistics_per_base: float
	logistics_bottleneck: str
	output_per_day: float
	average_traded_7d: float
	market_capacity_base: float
	ship_capex_per_base: float
	hq_costs: dict[str, dict[str, float]] # Added the HQ pricing dictionary

	def __lt__(self, other: Profit) -> bool:
		bases_a = self.area / 500
		p_a = (self.revenue['vwap'] - self.opex['vwap']) / bases_a
		c_a = self.capex['vwap'] / bases_a
		o_a = self.opex['vwap'] / bases_a
		be_a = (c_a + 3 * o_a) / p_a if p_a > 0 else 10000 - p_a
		
		bases_b = other.area / 500
		p_b = (other.revenue['vwap'] - other.opex['vwap']) / bases_b
		c_b = other.capex['vwap'] / bases_b
		o_b = other.opex['vwap'] / bases_b
		be_b = (c_b + 3 * o_b) / p_b if p_b > 0 else 10000 - p_b
		return be_a < be_b

@dataclasses.dataclass(eq=False, frozen=True, slots=True)
class MatPrice:
	ticker: str
	amount: int
	vwap_7d: float
	bid: float | None
	ask: float | None

if __name__ == '__main__':
	main()