select-deco/process_data.py
2023-12-01 15:50:27 +08:00

127 lines
3.9 KiB
Python

import json
import csv
from tqdm import tqdm
from scipy.special import comb
from operator import attrgetter
from typing import Any
from itertools import combinations
config = json.load(open('config.json', 'r', encoding='utf-8'))
def calc_weight(effect: str, level: int) -> float:
weight = 0
for tier in config['weight_tier']:
if level >= config['weight_tier'][tier]['gte'] and level < config['weight_tier'][tier]['lt']:
weight = config['weight_tier'][tier]['weight']
if effect in config['weight_prefer']:
weight *= config['weight_prefer'][effect]
return weight
class Item:
name: str
class_name: str
effect_list: dict[str, int]
total_level: int
def __init__(self, name: str, class_name: str) -> None:
self.name = name
self.class_name = class_name
self.effect_list = {}
self.total_level = 0
class Combo:
selected_item: list[Item]
total_effect: dict[str, int]
total_weight: float
def __init__(self, selected_item: list[Item]) -> None:
self.selected_item = selected_item
effect_level_map = {}
self.total_weight = 0
orig_item: list[Item] = []
attach_item: list[Item] = []
for item in selected_item:
if '' in item.name:
attach_item.append(item)
else:
orig_item.append(item)
for (effect, level) in item.effect_list.items():
if effect not in effect_level_map:
effect_level_map[effect] = level
else:
effect_level_map[effect] += level
self.total_effect = effect_level_map
if len(effect_level_map) > 5:
return
for item in orig_item:
for it in attach_item:
if item.name in it.name:
return
for (effect, level) in effect_level_map.items():
self.total_weight += calc_weight(effect, level)
class DictEncoder(json.JSONEncoder):
def default(self, o: Any) -> Any:
return o.__dict__
item_list: list[Item] = []
combo_list: list[Combo] = []
for src in config['classes']:
data_file = f'data/{src}.csv'
with open(data_file, 'r', encoding='utf-8') as csvfile:
reader = csv.reader(csvfile)
next(reader, None)
for row in reader:
item = Item(name=row[0], class_name=src)
effects = row[1:]
for effect in effects:
if '+' not in effect:
continue
effect_name = effect.split('+')[0]
effect_level = int(effect.split('+')[1])
item.effect_list[effect_name] = effect_level
item.total_level += effect_level
item_list.append(item)
filtered_item = []
for item in item_list:
if len(item.effect_list) == 0:
continue
if item.total_level <= config['min_item_level']:
continue
if item.class_name in config['exclude_class']:
if item.name in config['include_item']:
filtered_item.append(item)
continue
filtered_item.append(item)
filtered_item.sort(key=attrgetter('total_level'), reverse=True)
print(f'total {len(filtered_item)} items')
json.dump(filtered_item, open('filtered_item.json', 'w', encoding='utf-8'),
cls=DictEncoder, ensure_ascii=False, indent=2)
item_combo = combinations(filtered_item, config['select_count'])
tbar = tqdm(total=comb(len(filtered_item), config['select_count']))
for select_list in item_combo:
combo = Combo(selected_item=list(select_list))
if len(combo.total_effect) <= 5:
combo_list.append(combo)
tbar.update()
tbar.close()
result = sorted(combo_list, key=attrgetter('total_weight'), reverse=True)
json.dump(result[:10], open('result.json', 'w', encoding='utf-8'),
cls=DictEncoder, ensure_ascii=False, indent=2)
json.dump(result, open('result.full.json', 'w', encoding='utf-8'),
cls=DictEncoder, ensure_ascii=False, indent=2)