prunstats/py/prepare_from_sqlite.py

from __future__ import annotations

import collections
import sqlite3
import sys
import typing

import prepare

def main() -> None:
	month, sqlite_path, run_id_s = sys.argv[1:]
	run_id = int(run_id_s)

	db = sqlite3.connect(sqlite_path)
	db.row_factory = sqlite3.Row

	cur = db.execute('SELECT started_at, finished_at, type, range, materials_failed FROM runs WHERE id = ?', (run_id,))
	(run,) = cur.fetchall()
	print(', '.join(f'{k}: {run[k]}' for k in run.keys())) # noqa: SIM118
	assert run['type'] == 'PRODUCTION' and run['range'] == 'DAYS_30' and run['materials_failed']  == 0

	with open(f'rawData/{month}-prices.json', 'r') as f:
		prices = prepare.get_prices(f)
	prod_data, company_data = get_prod_and_company_data(db, run_id, prices)

	with open(f'rawData/{month}.csv', 'r', newline='') as f:
		data = prepare.read_data(f)
	prepare.write_data(month, data, prod_data, company_data)
	prepare.check_missing_tickers(prod_data)

def get_prod_and_company_data(db: sqlite3.Connection, run_id: int, prices: typing.Mapping[str, float]
		) -> tuple[typing.Mapping[str, prepare.ProdData], typing.Mapping[str, typing.Any]]:
	cur = db.execute('''
			SELECT ticker, entity_id, score, rank FROM leaderboard_scores
			JOIN materials ON materials.material_pk = leaderboard_scores.material_rowid
			JOIN entities ON entities.id = leaderboard_scores.entity_rowid
			WHERE run_id = ?''', (run_id,))

	individual: dict[str, dict[str, prepare.CompanyTickerData]] = collections.defaultdict(dict)
	company_totals: dict[str, prepare.CompanyTotals] = collections.defaultdict(lambda: {'volume': 0.0})
	universe_total_score: dict[str, float] = collections.defaultdict(int)

	while row := cur.fetchone():
		ticker = row['ticker']
		price = prices[ticker]
		universe_total_score[ticker] += row['score']
		amount = row['score'] / 30
		volume = amount * price
		individual[row['entity_id']][ticker] = {
			'amount': amount,
			'volume': volume,
			'rank': row['rank'],
		}
		company_totals[row['entity_id']]['volume'] += volume

	prod: dict[str, prepare.ProdData] = {} # TODO: universe / 30
	for ticker, total_score in universe_total_score.items():
		amount = total_score / 30
		prod[ticker] = {'amount': amount, 'volume': amount * prices[ticker]}

	company_data = {'totals': prepare.add_company_ranks(company_totals), 'individual': dict(individual)}
	return prod, company_data

if __name__ == '__main__':
	main()