from math import ceil, floor
from typing import Union

from ..Class.newLocationClass import KH2Location
from ..List.LvupStats import DreamWeaponOffsets
from ..List.NewLocationList import Locations
from ..List.configDict import (
    itemType,
    locationCategory,
    itemRarity,
    locationDepth,
    locationType,
    itemBias,
    itemDifficulty,
)
from ..List.location.graph import START_NODE
from .RandomizerSettings import RandomizerSettings


class SimplifiedWeightDistributions:
    def __init__(self, max_depth: int):
        self.weighting_function: dict[itemBias, list[int]] = {}

        weight_list_maker = lambda function_index: [
            max(
                1,
                floor(
                    10.0
                    / pow(
                        (max_depth + 2 - x) / (max_depth / 3.0),
                        function_index * 0.5 + 1,
                    )
                ),
            )
            for x in range(1, max_depth + 2)
        ]

        self.weighting_function[itemBias.VERY_EARLY] = weight_list_maker(1)[::-1]
        self.weighting_function[itemBias.EARLY] = weight_list_maker(0)[::-1]
        self.weighting_function[itemBias.SLIGHTLY_EARLY] = [2] * floor(
            (max_depth + 1) / 2
        ) + [1] * ceil((max_depth + 1) / 2)
        self.weighting_function[itemBias.NOBIAS] = [1] * (max_depth + 1)
        self.weighting_function[itemBias.SLIGHTLY_LATE] = [1] * floor(
            (max_depth + 1) / 2
        ) + [2] * ceil((max_depth + 1) / 2)
        self.weighting_function[itemBias.LATE] = weight_list_maker(0)
        self.weighting_function[itemBias.VERY_LATE] = weight_list_maker(1)
        self.weighting_function[itemBias.SUPER_LATE] = weight_list_maker(2)
        self.weighting_function[itemBias.NIGHTMARE] = weight_list_maker(3)

    def get_rarity_weightings(self):
        return self.weighting_function


class WeightDistributions:

    def __init__(self, max_depth: int):
        self.weighting_function: dict[itemDifficulty, dict[itemRarity, list[int]]] = {
            itemDifficulty.NORMAL: {}
        }
        for rarity in itemRarity:
            self.weighting_function[itemDifficulty.NORMAL][itemRarity(rarity)] = [1] * (
                max_depth + 1
            )

        self.weighting_function[itemDifficulty.SLIGHTLY_HARD] = {
            itemRarity.COMMON: [1] * (max_depth + 1),
            itemRarity.UNCOMMON: [1] * (max_depth + 1),
            itemRarity.RARE: [1] * floor((max_depth + 1) / 2)
            + [2] * ceil((max_depth + 1) / 2),
            itemRarity.MYTHIC: [1] * floor((max_depth + 1) / 2)
            + [2] * ceil((max_depth + 1) / 2),
        }
        self.weighting_function[itemDifficulty.SLIGHTLY_EASY] = {
            itemRarity.COMMON: [1] * (max_depth + 1),
            itemRarity.UNCOMMON: [1] * (max_depth + 1),
            itemRarity.RARE: [2] * floor((max_depth + 1) / 2)
            + [1] * ceil((max_depth + 1) / 2),
            itemRarity.MYTHIC: [2] * floor((max_depth + 1) / 2)
            + [1] * ceil((max_depth + 1) / 2),
        }

        easier_difficulties = [itemDifficulty.SUPEREASY, itemDifficulty.EASY]
        for diff_index, difficulty in enumerate(easier_difficulties):
            diff_index_modded = (1 - diff_index) - 0.25
            self.weighting_function[difficulty] = {
                itemRarity.COMMON: [
                    max(
                        1,
                        floor(
                            10.0
                            / pow(x / (max_depth / 3.0), diff_index_modded * 0.5 + 1)
                        ),
                    )
                    for x in range(max_depth + 1, 0, -1)
                ],
                itemRarity.UNCOMMON: [
                    max(
                        1,
                        floor(
                            3.0
                            / pow(x / (max_depth / 3.0), diff_index_modded * 0.5 + 1)
                        ),
                    )
                    for x in range(max_depth + 1, 0, -1)
                ],
                itemRarity.RARE: [
                    max(
                        1,
                        floor(
                            3.0
                            / pow(
                                (max_depth + 2 - x) / (max_depth / 3.0),
                                diff_index_modded * 0.5 + 1,
                            )
                        ),
                    )
                    for x in range(max_depth + 1, 0, -1)
                ],
                itemRarity.MYTHIC: [
                    max(
                        1,
                        floor(
                            10.0
                            / pow(
                                (max_depth + 2 - x) / (max_depth / 3.0),
                                diff_index_modded * 0.5 + 1,
                            )
                        ),
                    )
                    for x in range(max_depth + 1, 0, -1)
                ],
            }

        harder_difficulties = [
            itemDifficulty.HARD,
            itemDifficulty.VERYHARD,
            itemDifficulty.INSANE,
            itemDifficulty.NIGHTMARE,
        ]
        for diff_index, difficulty in enumerate(harder_difficulties):
            diff_index_modded = diff_index - 0.25
            self.weighting_function[difficulty] = {
                itemRarity.COMMON: [
                    max(
                        1,
                        floor(
                            10.0
                            / pow(x / (max_depth / 3.0), diff_index_modded * 0.5 + 1)
                        ),
                    )
                    for x in range(1, max_depth + 2)
                ],
                itemRarity.UNCOMMON: [
                    max(
                        1,
                        floor(
                            3.0
                            / pow(x / (max_depth / 3.0), diff_index_modded * 0.5 + 1)
                        ),
                    )
                    for x in range(1, max_depth + 2)
                ],
                itemRarity.RARE: [
                    max(
                        1,
                        floor(
                            3.0
                            / pow(
                                (max_depth + 2 - x) / (max_depth / 3.0),
                                diff_index_modded * 0.5 + 1,
                            )
                        ),
                    )
                    for x in range(1, max_depth + 2)
                ],
                itemRarity.MYTHIC: [
                    max(
                        1,
                        floor(
                            10.0
                            / pow(
                                (max_depth + 2 - x) / (max_depth / 3.0),
                                diff_index_modded * 0.5 + 1,
                            )
                        ),
                    )
                    for x in range(1, max_depth + 2)
                ],
            }

        # for diff_index,difficulty in enumerate(["Super Easy","Easy","Hard","Very Hard","Insane","Nightmare"]):
        #     print(self.weighting_function[difficulty][itemRarity.COMMON])
        #     print(self.weighting_function[difficulty][itemRarity.UNCOMMON])
        #     print(self.weighting_function[difficulty][itemRarity.RARE])
        #     print(self.weighting_function[difficulty][itemRarity.MYTHIC])
        #     print("--------------------")

    def get_rarity_weighting(self, difficulty: itemDifficulty):
        if difficulty not in self.weighting_function:
            raise ValueError(f"Unknown item placement difficulty {difficulty}")
        return self.weighting_function[difficulty]


class LocationWeights:

    def __init__(
        self,
        settings: RandomizerSettings,
        locations: Locations,
        reverse_locations: Locations,
    ):
        self.regular_rando = settings.regular_rando
        self.reverse_rando = settings.reverse_rando
        self.split_levels = settings.split_levels
        self.level_offsets = DreamWeaponOffsets()
        self.max_level = settings.max_level_checks
        self.level_depths: dict[int, int] = {}

        # Regular
        # -------------------------------
        self.location_depths: dict[KH2Location, int] = {}
        hops = locations.location_graph.get_hops(START_NODE)
        location_type_maxes: dict[locationType, int] = {}
        max_hops = (
            17  # hard coding max hops to unify depths between reverse and regular rando
        )
        # for hop in hops:
        #     max_hops = max(max_hops,hop[1])
        for node, distance in hops:
            for loc in locations.locations_for_node(node):
                primary_type = loc.LocationTypes[0]
                if primary_type not in location_type_maxes:
                    location_type_maxes[primary_type] = distance
                else:
                    location_type_maxes[primary_type] = max(
                        distance, location_type_maxes[primary_type]
                    )

        for node, distance in hops:
            for loc in locations.locations_for_node(node):
                primary_type = loc.LocationTypes[0]
                if location_type_maxes[primary_type] != 0:
                    scaled_depth = floor(
                        (distance * 1.0 / location_type_maxes[primary_type]) * max_hops
                    )
                else:
                    scaled_depth = distance
                self.location_depths[loc] = scaled_depth
                if loc.LocationCategory is locationCategory.LEVEL:
                    self.level_depths[loc.LocationId] = scaled_depth

        # Reverse
        # -------------------------------
        self.reverse_location_depths: dict[KH2Location, int] = {}
        hops = reverse_locations.location_graph.get_hops(START_NODE)
        reverse_location_type_maxes: dict[locationType, int] = {}
        # max_hops = 0
        # for hop in hops:
        #     max_hops = max(max_hops,hop[1])
        for node, distance in hops:
            for loc in reverse_locations.locations_for_node(node):
                primary_type = loc.LocationTypes[0]
                if primary_type not in reverse_location_type_maxes:
                    reverse_location_type_maxes[primary_type] = distance
                else:
                    reverse_location_type_maxes[primary_type] = max(
                        distance, reverse_location_type_maxes[primary_type]
                    )

        for node, distance in hops:
            for loc in reverse_locations.locations_for_node(node):
                primary_type = loc.LocationTypes[0]
                if reverse_location_type_maxes[primary_type] != 0:
                    scaled_depth = floor(
                        (distance * 1.0 / reverse_location_type_maxes[primary_type])
                        * max_hops
                    )
                else:
                    scaled_depth = distance
                self.reverse_location_depths[loc] = scaled_depth

        self.base_weights = SimplifiedWeightDistributions(
            max_hops
        ).get_rarity_weightings()
        self.weights = {}
        # default to normal weighting
        for item_type in itemType:
            self.weights[item_type] = self.base_weights[itemBias.NOBIAS]

        self.weights[itemType.FORM] = self.base_weights[settings.form_weights]
        self.weights[itemType.STORYUNLOCK] = self.base_weights[settings.unlock_weights]

        self.weights[itemType.FIRE] = self.base_weights[settings.magic_weights]
        self.weights[itemType.BLIZZARD] = self.base_weights[settings.magic_weights]
        self.weights[itemType.THUNDER] = self.base_weights[settings.magic_weights]
        self.weights[itemType.CURE] = self.base_weights[settings.magic_weights]
        self.weights[itemType.MAGNET] = self.base_weights[settings.magic_weights]
        self.weights[itemType.REFLECT] = self.base_weights[settings.magic_weights]

        self.weights[itemType.TORN_PAGE] = self.base_weights[settings.page_weights]

        self.weights[itemType.SUMMON] = self.base_weights[settings.summon_weights]

        self.weights[itemType.REPORT] = self.base_weights[settings.report_weights]

        self.weights[itemType.PROOF_OF_PEACE] = self.base_weights[
            settings.proof_weights
        ]
        self.weights[itemType.PROOF_OF_CONNECTION] = self.base_weights[
            settings.proof_weights
        ]
        self.weights[itemType.PROOF_OF_NONEXISTENCE] = self.base_weights[
            settings.proof_weights
        ]

        self.weights[itemType.PROMISE_CHARM] = self.base_weights[
            settings.promise_charm_weights
        ]

    def get_depth(self, location: KH2Location) -> Union[int, tuple[int, int]]:
        """
        Returns either a single depth (for levels or regular or reverse rando) or a tuple of
        (regular depth, reverse depth) if both regular and reverse rando.
        """
        if location.LocationCategory is locationCategory.LEVEL:
            loc_id = location.LocationId
            if self.split_levels:
                sword_depth = self.level_depths[loc_id]
                shield_depth = self.level_depths[
                    self.level_offsets.get_shield_level(self.max_level, loc_id)
                ]
                staff_depth = self.level_depths[
                    self.level_offsets.get_staff_level(self.max_level, loc_id)
                ]
                return (sword_depth + shield_depth + staff_depth) // 3
            else:
                return self.level_depths[loc_id]
        else:
            if self.regular_rando and self.reverse_rando:
                return (
                    self.location_depths[location],
                    self.reverse_location_depths[location],
                )
            elif self.regular_rando:
                return self.location_depths[location]
            else:
                return self.reverse_location_depths[location]

    def get_weight(self, rarity: itemRarity, location: KH2Location) -> int:
        raise AssertionError("Deprecated weighting function")
        """ Returns the weight that should be used for an item of the given rarity at the given location. """
        depth_or_depths = self.get_depth(location)
        rarity_weights = self.weights[rarity]
        if isinstance(depth_or_depths, int):
            return rarity_weights[depth_or_depths]
        elif isinstance(depth_or_depths, tuple):
            regular_weight = rarity_weights[depth_or_depths[0]]
            reverse_weight = rarity_weights[depth_or_depths[1]]
            return (regular_weight + reverse_weight) // 2

    def get_weight(self, item_type: itemType, location: KH2Location) -> int:
        """Returns the weight that should be used for an item of the given type at the given location."""
        depth_or_depths = self.get_depth(location)
        rarity_weights = self.weights[item_type]
        if isinstance(depth_or_depths, int):
            return rarity_weights[depth_or_depths]
        elif isinstance(depth_or_depths, tuple):
            regular_weight = rarity_weights[depth_or_depths[0]]
            reverse_weight = rarity_weights[depth_or_depths[1]]
            return (regular_weight + reverse_weight) // 2
