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)