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

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[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}

	target_days = float('inf')
	for planet in planets:
		vol_per_day = weight_per_day = 0
		for consumption in planet.net_consumption:
			ticker = consumption['MaterialTicker']
			vol_per_day += materials[ticker]['Volume'] * consumption['net_consumption']
			weight_per_day += materials[ticker]['Weight'] * consumption['net_consumption']
			days = planet.inventory.get(ticker, 0) / consumption['net_consumption']
			if days < target_days:
				target_days = days

		print(planet.name, f'consumes {vol_per_day:.1f}㎥, {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, weight_used, volume_used = planet.buy_for_target(materials, target_days)
			iteration_weight += weight_used
			iteration_volume += volume_used
			if iteration_weight > args.weight or iteration_volume > args.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)}㎥') # pyright: ignore[reportPossiblyUnboundVariable]

	for planet in planets:
		print('\n' + cyan(planet.name))
		for consumption in planet.net_consumption:
			ticker = consumption['MaterialTicker']
			avail = planet.inventory.get(ticker, 0)
			daily_consumption = consumption['net_consumption']
			days = avail / daily_consumption
			print(f'{ticker:>3}: {avail:5d} ({daily_consumption:8.2f}/d) {days:4.1f} d', end='')
			if need := optimal[planet.name].get(ticker): # pyright: ignore[reportOptionalMemberAccess]
				print(f' | {need:8.1f}')
			else:
				print()

	combined_buy: dict[str, int] = collections.defaultdict(int)
	for buy in optimal.values():
		for ticker, amount in buy.items():
			combined_buy[ticker] += amount
	print(cyan('\nbuy:\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 combined_buy.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 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[Inventory]
	OrderProduction: list[Amount]

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

	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']}
		producing = {item['MaterialTicker']: item for item in fio_burn['OrderProduction']}
		self.net_consumption = []
		for c in fio_burn['OrderConsumption'] + fio_burn['WorkforceConsumption']:
			net = c['DailyAmount']
			if production := producing.get(c['MaterialTicker']):
				net -= production['DailyAmount']
				if net < 0:
					continue
			c['net_consumption'] = net
			self.net_consumption.append(c)

	def buy_for_target(self, materials: dict[str, Material], target_days: float) -> tuple[dict[str, int], float, float]:
		weight_used = volume_used = 0
		buy: dict[str, int] = {}
		for consumption in self.net_consumption:
			ticker = consumption['MaterialTicker']
			avail = self.inventory.get(ticker, 0)
			daily_consumption = consumption['net_consumption']
			days = avail / daily_consumption
			if days < target_days:
				buy[ticker] = math.ceil((target_days - days) * daily_consumption)
				weight_used += buy[ticker] * materials[ticker]['Weight']
				volume_used += buy[ticker] * materials[ticker]['Volume']
		return buy, weight_used, volume_used

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

class Inventory(typing.TypedDict):
	MaterialTicker: str
	MaterialAmount: int

class Storage(typing.TypedDict):
	Name: str
	StorageItems: list[Inventory]
	Type: typing.Literal['STORE'] | typing.Literal['WAREHOUSE_STORE'] | typing.Literal['FTL_FUEL_STORE'] | typing.Literal['STL_FUEL_STORE'] |  typing.Literal['SHIP_STORE']

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

if __name__ == '__main__':
	main()