from typing import Union
from pathlib import Path
from typing import Dict
from typing import Set
import logging
import json


def load_forces_from_json(json_path: Union[str, Path]) -> Dict[str, Dict[str, Union[Set, Dict]]]:
    """
    Load forced labels, changes, and non-changes from a json file and return them
    in a dictionary that can be passed through to JointModel.get_harmony(...) as kwargs.

    Parameters
    ----------
    json_path : Union[str, Path]
        A reference to a json file containing forced labels, changes, and non-changes,
        to be used for JointModel.get_harmony(...). It may have the following keys:
            - "forced_chord_changes": A list of integers at which the chord must change.
            - "forced_chord_non_changes": A list of integers at which the chord cannot change.
            - "forced_key_changes": A list of integers at which the key must change.
            - "forced_key_non_changes": A list of integers at which the key cannot change.
            - "forced_chords": A dictionary mapping the string form of a tuple in the form
                               (start, end) to a chord_id, saying that the input indexes on
                               the range start (inclusive) to end (exclusive) must be output
                               as the given chord_id.
            - "forced_keys": Same as forced_chords, but for keys.

    Returns
    -------
    forces_kwargs: Dict[str, Dict[str, Union[Set, Dict]]]
        A nested dictionary containing the loaded keyword arguments for each input.
        The outer-most keys should reference a specific input by string name,
        or be the keyword "default", in which case the loaded kwargs will be used for all
        input pieces not otherwise matched by string name.
        In the inner dictionaries, keyword arguments have been loaded (with the correct types)
        from the json file that can be passed directly as kwargs to JointModel.get_harmony(...)
        for that particular piece.
    """

    def load_forces_from_nested_json(raw_data: Dict) -> Dict[str, Union[Set, Dict]]:
        """
        Load an inner forces_kwargs dict from a nested json forces_kwargs dict data.

        Parameters
        ----------
        raw_data : Dict
            The inner nested dictionary from which we will load the kwargs.
            See load_forces_from_json for details.

        Returns
        -------
        Dict[str, Union[Set, Dict]]
            The kwargs for a single piece, unnested.
        """
        forces_kwargs = dict()

        for key in [
            "forced_chord_changes",
            "forced_chord_non_changes",
            "forced_key_changes",
            "forced_key_non_changes",
        ]:
            if key in raw_data:
                forces_kwargs[key] = set(raw_data[key])

        for key in ["forced_chords", "forced_keys"]:
            if key in raw_data:
                forces_kwargs[key] = {
                    tuple(map(int, range_tuple_str[1:-1].split(","))): label_id
                    for range_tuple_str, label_id in raw_data[key].items()
                }

        for key in raw_data:
            if key not in [
                "forced_chord_changes",
                "forced_chord_non_changes",
                "forced_key_changes",
                "forced_key_non_changes",
                "forced_chords",
                "forced_keys",
            ]:
                logging.warning(
                    "--forces-json inner key not recognized: %s. Ignoring that key.", key
                )

        logging.info("Forces:" if len(forces_kwargs) > 0 else "Forces: None")
        for key, item in sorted(forces_kwargs.items()):
            if type(item) == dict:
                logging.info("    %s:", key)
                for inner_key, inner_item in sorted(item.items()):
                    logging.info("        %s = %s", inner_key, inner_item)
            else:
                logging.info("    %s = %s", key, item)

        return forces_kwargs

    with open(json_path, "r") as json_file:
        raw_data = json.load(json_file)

    if (
        "forced_chord_changes" in raw_data
        or "forced_chord_non_changes" in raw_data
        or "forced_key_changes" in raw_data
        or "forced_key_non_changes" in raw_data
        or "forced_chords" in raw_data
        or "forced_keys" in raw_data
    ):
        logging.info(
            "Given --json-forces not a nested, piece-specific mapping. Treating as default for "
            "all inputs."
        )
        raw_data = {"default": raw_data}

    all_forces_kwargs = {}

    for key, nested_raw_data in raw_data.items():
        logging.info("Loading forces for %s", key)
        all_forces_kwargs[key] = load_forces_from_nested_json(nested_raw_data)

    return all_forces_kwargs