pruncalc/supply.py

from __future__ import annotations

import collections
import dataclasses
import json
import math
import typing

import tap

import cache
from config import config
import market

class Args(tap.Tap):
	planets: tuple[str, ...]
	weight: float
	volume: float
	include_ship: tuple[str, ...] = ()

	def configure(self) -> None:
		self.add_argument('planets', nargs='+', metavar='planet') # take planets as positional args instead of flag

def main() -> None:
	args = Args().parse_args()

	planets = [Planet(fio_burn) for fio_burn in get_fio_burn(args.planets)]
	if args.include_ship:
		stores: typing.Sequence[market.Storage] = cache.get('https://rest.fnar.net/storage/' + config.username,
				headers={'Authorization': config.fio_api_key})
		for ship in args.include_ship:
			ship_name, planet_name = ship.casefold().split('=')
			for store in stores:
				if store['Type'] == 'SHIP_STORE' and store['Name'].casefold() == ship_name:
					break
			else:
				raise Exception(f'ship storage {ship_name} not found')
			(planet,) = (p for p in planets if p.name.casefold() == planet_name)
			for item in store['StorageItems']:
				planet.inventory[item['MaterialTicker']] = planet.inventory.get(item['MaterialTicker'], 0) + item['MaterialAmount']

	raw_materials: typing.Sequence[Material] = cache.get('https://rest.fnar.net/material/allmaterials')
	materials = {mat['Ticker']: mat for mat in raw_materials}

	optimal, from_cx = calculate_optimal(planets, materials, args.weight, args.volume)

	weight_used = volume_used = 0
	for mat, amount in from_cx.items():
		weight_used += amount * materials[mat]['Weight']
		volume_used += amount * materials[mat]['Volume']
	for planet in planets:
		unload = optimal[planet.name]
		for mat, amount in unload.items():
			weight_used -= amount * materials[mat]['Weight']
			volume_used -= amount * materials[mat]['Volume']
		for mat in planet.exporting:
			amount = planet.inventory.get(mat, 0)
			weight_used += amount * materials[mat]['Weight']
			volume_used += amount * materials[mat]['Volume']
		if weight_used > args.weight:
			print(yellow('warning'), f'need additional {weight_used - args.weight:.1f}t to pick up exports at {planet.name}')
		if volume_used > args.volume:
			print(yellow('warning'), f'need additional {volume_used - args.volume:.1f}m³ to pick up exports at {planet.name}')
	print(cyan('\nload at CX:\n') + json.dumps({
		'actions': [
			{'name': 'BuyItems', 'type': 'CX Buy', 'group': 'A1', 'exchange': 'IC1',
				'priceLimits': {}, 'buyPartial': False, 'useCXInv': True},
			{'type': 'MTRA', 'name': 'TransferAction', 'group': 'A1',
				'origin': 'Hortus Station Warehouse', 'dest': 'Configure on Execution'},
		],
		'global': {'name': 'supply ' + ' '.join(args.planets)},
		'groups': [{
			'type': 'Manual', 'name': 'A1', 'materials': {mat: amount for mat, amount in from_cx.items()}
		}],
	}))
	for planet in planets:
		buy = optimal[planet.name]
		print(cyan(f'unload {planet.name}:\n') + json.dumps({
			'actions': [
				{'type': 'MTRA', 'name': 'TransferAction', 'group': 'A1',
					'origin': 'Configure on Execution', 'dest': planet.name + ' Base'},
			],
			'global': {'name': 'unload ' + planet.name},
			'groups': [{
				'type': 'Manual', 'name': 'A1', 'materials': {mat: amount for mat, amount in buy.items()}
			}],
		}))

def get_fio_burn(planet_names: typing.Sequence[str]) -> typing.Iterator[FIOBurn]:
	planets: list[FIOBurn] = cache.get('https://rest.fnar.net/fioweb/burn/user/' + config.username,
			headers={'Authorization': config.fio_api_key})
	for name in planet_names:
		name = name.casefold()
		for planet_data in planets:
			if name in (planet_data['PlanetName'].casefold(), planet_data['PlanetNaturalId'].casefold()):
				assert planet_data['Error'] is None
				yield planet_data
				break
		else:
			raise ValueError(name + ' not found')

def calculate_optimal(planets: typing.Sequence[Planet], materials: typing.Mapping[str, Material], max_weight: float,
			max_volume: float) -> tuple[typing.Mapping[str, typing.Mapping[str, int]], typing.Mapping[str, int]]:
	target_days = float('inf')
	for planet in planets:
		vol_per_day = weight_per_day = 0
		for ticker, consumption in planet.net_consumption.items():
			if consumption <= 0:
				continue
			vol_per_day += materials[ticker]['Volume'] * consumption
			weight_per_day += materials[ticker]['Weight'] * consumption
			days = planet.inventory.get(ticker, 0) / consumption
			if days < target_days:
				target_days = days

		print(planet.name, f'consumes {vol_per_day:.1f}m³, {weight_per_day:.1f}t per day')

	optimal: dict[str, dict[str, int]] = None # pyright: ignore[reportAssignmentType]
	total_weight_used: float = None # pyright: ignore[reportAssignmentType]
	total_volume_used: float = None # pyright: ignore[reportAssignmentType]
	target_days = round(target_days + 0.05, 1)
	load_more = True
	while load_more:
		buys: dict[str, dict[str, int]] = {}
		iteration_weight = iteration_volume = 0
		for planet in planets:
			buy = planet.supply_for_days(target_days)
			weight_used, volume_used = shipping_used(materials, config.supply_config(planet.name).ignore_materials, buy)
			iteration_weight += weight_used
			iteration_volume += volume_used
			if iteration_weight > max_weight or iteration_volume > max_volume:
				load_more = False
				break
			buys[planet.name] = buy
		if load_more:
			optimal = buys
			total_weight_used = iteration_weight
			total_volume_used = iteration_volume
		target_days += 0.1
	print('supply for', round(target_days, 1), 'days,', end=' ')
	print(f'consuming {round(total_weight_used, 1)}t and {round(total_volume_used, 1)}m³') # pyright: ignore[reportPossiblyUnboundVariable]

	raw_prices: typing.Mapping[str, market.RawPrice] = {p['MaterialTicker']: p
			for p in cache.get('https://refined-prun.github.io/refined-prices/all.json') if p['ExchangeCode'] == 'IC1'}

	warehouse = warehouse_inventory()
	from_cx: dict[str, int] = collections.defaultdict(int)
	total_cost = 0
	for i, planet in enumerate(planets):
		print('\n' + cyan(planet.name))
		supply_config = config.supply_config(planet.name)
		planet_buy = optimal[planet.name]
		for ticker, consumption in planet.net_consumption.items():
			if consumption <= 0:
				continue
			avail = planet.inventory.get(ticker, 0)
			days = avail / consumption
			print(f'{ticker:>3}: {avail:5d} ({consumption:8.2f}/d) {days:5.1f} d', end='')
			if need := planet_buy.get(ticker): # pyright: ignore[reportOptionalMemberAccess]
				if ticker in supply_config.ignore_materials:
					print(f' | {need:5.0f} (ignored)')
				else:
					print(f' | {need:5.0f}', end='')
					sources = []
					for exporter in planets[:i]:
						if ticker in exporter.exporting and (avail := exporter.inventory.get(ticker, 0)):
							need -= min(need, avail)
							exporter.inventory[ticker] -= max(avail - need, 0)
							sources.append(f'{exporter.name}: {avail}')
					if need:
						from_cx[ticker] += need # count from_cx before subtracting warehouse
					if avail := warehouse.get(ticker, 0):
						need -= min(need, avail)
						warehouse[ticker] -= max(avail - need, 0)
						sources.append(f'WH: {avail}')
					if (ppu := raw_prices[ticker]['Ask']) is not None:
						cost = ppu * need
						print(f' (${cost:6.0f}) ' + ', '.join(sources))
						total_cost += cost
					else:
						print(yellow(' no supply ') + ', '.join(sources))
			else:
				print()
	print(f'\ntotal cost: {total_cost:,}')
	return optimal, from_cx

def shipping_used(materials: typing.Mapping[str, Material], ignore: typing.Collection[str], counts: dict[str, int]) -> tuple[float, float]:
	weight = volume = 0
	for ticker, amount in counts.items():
		if ticker in ignore:
			continue
		weight += amount * materials[ticker]['Weight']
		volume += amount * materials[ticker]['Volume']
	return weight, volume

def warehouse_inventory() -> dict[str, int]:
	warehouses: typing.Sequence[market.Warehouse] = cache.get('https://rest.fnar.net/sites/warehouses/' + config.username,
			headers={'Authorization': config.fio_api_key})
	for warehouse in warehouses:
		if warehouse['LocationNaturalId'] == 'HRT':
			storage: market.Storage = cache.get(f'https://rest.fnar.net/storage/{config.username}/{warehouse["StoreId"]}',
				headers={'Authorization': config.fio_api_key})
			assert storage['Type'] == 'WAREHOUSE_STORE'
			return {item['MaterialTicker']: item['MaterialAmount'] for item in storage['StorageItems']}
	raise Exception("couldn't find HRT warehouse")

def yellow(text: str) -> str:
	return '\033[33m' + text + '\033[0m'

def cyan(text: str) -> str:
	return '\033[36m' + text + '\033[0m'

class FIOBurn(typing.TypedDict):
	PlanetName: str
	PlanetNaturalId: str
	Error: typing.Any
	OrderConsumption: list[Amount]
	WorkforceConsumption: list[Amount]
	Inventory: list[market.StorageItem]
	OrderProduction: list[Amount]

class Amount(typing.TypedDict):
	MaterialTicker: str
	DailyAmount: float

@dataclasses.dataclass(init=False, eq=False, slots=True)
class Planet:
	name: str
	inventory: dict[str, int]
	net_consumption: dict[str, float]
	exporting: typing.Set[str]

	def __init__(self, fio_burn: FIOBurn) -> None:
		self.name = fio_burn['PlanetName'] or fio_burn['PlanetNaturalId']
		self.inventory = {item['MaterialTicker']: item['MaterialAmount'] for item in fio_burn['Inventory']}

		self.net_consumption = {}
		for item in fio_burn['OrderConsumption'] + fio_burn['WorkforceConsumption']:
			ticker = item['MaterialTicker']
			self.net_consumption[ticker] = self.net_consumption.get(ticker, 0) + item['DailyAmount']
		for item in fio_burn['OrderProduction']:
			ticker = item['MaterialTicker']
			self.net_consumption[ticker] = self.net_consumption.get(ticker, 0) - item['DailyAmount']

		# producing more than consumption
		self.exporting = set()
		for item in fio_burn['OrderProduction']:
			if item['MaterialTicker'] not in self.net_consumption:
				self.exporting.add(item['MaterialTicker'])

	def supply_for_days(self, target_days: float) -> dict[str, int]:
		buy: dict[str, int] = {}
		for ticker, consumption in self.net_consumption.items():
			if consumption <= 0:
				continue
			avail = self.inventory.get(ticker, 0)
			days = avail / consumption
			to_buy = 0
			if days < target_days:
				to_buy = math.ceil((target_days - days) * consumption)
			if avail + to_buy < 2:
				to_buy = 2 - avail # always stock at least 2 of everything
			if to_buy > 0:
				buy[ticker] = to_buy
		return buy

class Material(typing.TypedDict):
	Ticker: str
	Weight: float
	Volume: float

if __name__ == '__main__':
	main()