from typing import Iterator, Tuple, Dict, List
import numpy as np
from nptyping import NDArray
from omegaconf import DictConfig
from ..objective import Objective, DemoMonotone, DemoMonotoneSkewed, \
                       DemoNonMonotone, FacilityLocation, BudgetAllocation
from .. import dataset_utils


def compute_B(rng: np.random.Generator,
              n: int, B_range: Tuple[int, int]) -> NDArray[int]:
    """
    Compute the upper-bound array B.
    :param rng: numpy random generator instance
    :param n: size of the ground set
    :param B_range: inclusive range for the values of each entry of B
    """
    low, high = B_range

    return rng.integers(low=low, high=high, size=(n, ), endpoint=True)


OBJ_MAP = {
    'demo_monotone': lambda **kwargs: load_demo_monotone(**kwargs),
    'demo_monotone_skewed': lambda **kwargs: load_demo_monotone_skewed(**kwargs),
    'demo_non_monotone': lambda **kwargs: load_demo_non_monotone(**kwargs),
    'facility_location': lambda **kwargs: load_facility_location(**kwargs),
    'budget_allocation': lambda **kwargs: load_budget_allocation(**kwargs),
}


def load_demo_monotone(rng: np.random.Generator,
                       params,
                       **kwargs) -> List[Tuple[Objective, int]]:
    """
    Generate a random modular, monotone function
    :param rng: numpy random generator instance
    :param params: 'params.demo_monotone' dictionary entry in conf/config.yaml
    """
    nBr: List[Tuple[int, Tuple[int, int], int]] = params.benchmark.nBr

    fr: List[Tuple[Objective, int]] = [None] * len(nBr)
    for i, (n, B_range, r) in enumerate(nBr):
        B = compute_B(rng, n, B_range)
        fr[i] = (DemoMonotone(rng, n=n, B=B, B_range=B_range), r)

    return fr


def load_demo_monotone_skewed(rng: np.random.Generator,
                              params,
                              **kwargs) -> List[Tuple[Objective, int]]:
    """
    Generate a random skewed modular, monotone function
    :param rng: numpy random generator instance
    :param params: 'params.demo_monotone' dictionary entry in conf/config.yaml
    """
    nBr: List[Tuple[int, Tuple[int, int], int]] = params.benchmark.nBr

    fr: List[Tuple[Objective, int]] = [None] * len(nBr)
    for i, (n, B_range, r) in enumerate(nBr):
        B = compute_B(rng, n, B_range)
        fr[i] = (DemoMonotoneSkewed(rng, n=n, B=B, B_range=B_range), r)

    return fr


def load_demo_non_monotone(rng: np.random.Generator,
                           params,
                           **kwargs) -> List[Tuple[Objective, int]]:
    """
    Generate a random modular, non_monotone function
    :param rng: numpy random generator instance
    :param params: 'params.demo_non_monotone' dictionary entry in conf/config.yaml
    """
    nBr: List[Tuple[int, Tuple[int, int], int]] = params.benchmark.nBr

    fr: List[Tuple[Objective, int]] = [None] * len(nBr)
    for i, (n, B_range, r) in enumerate(nBr):
        B = compute_B(rng, n, B_range)
        fr[i] = (DemoNonMonotone(rng, n=n, B=B, B_range=B_range), r)

    return fr


def load_facility_location(rng: np.random.Generator,
                           dataset_dir: str,
                           params,
                           **kwargs) -> List[Tuple[Objective, int]]:
    """
    Generate a random integer-lattice submodular, monotone function that models
    the Facility Location problem.
    :param rng: numpy random generator instance
    :param dataset_dir: datasets main directory
    :param params: 'params.demo_facility_location' dictionary entry in conf/config.yaml
    """
    print(f'Loading Movielens 100k...')
    G = dataset_utils.import_movielens_100k(dataset_dir)
    print(f'...Movielens 100k successfully loaded')
    br: List[Tuple[int, int]] = params.benchmark.br
    
    fr: List[Tuple[Objective, int]] = [None] * len(br)
    for i, (b, r) in enumerate(br):
        fr[i] = (FacilityLocation(G=G, b=b), r)

    return fr


def load_budget_allocation(rng: np.random.Generator,
                           dataset_dir: str,
                           params,
                           **kwargs) -> Iterator[Tuple[Objective, int]]:
    """
    Generate a random integer-lattice DR-submodular, monotone function that models
    the Budget Allocation problem.
    :param rng: numpy random generator instance
    :param dataset_dir: datasets main directory
    :param params: 'params.demo_facility_location' dictionary entry in conf/config.yaml
    """
    # keep track of the vector B generated by the corresponding (n, b_low, b_high)
    n_B_range_to_B_map: Dict[Tuple[int, int, int], NDArray[int]] = dict()

    trim_graph = dataset_utils.import_wikilens_ratings(rng, dataset_dir)
    nBr: List[Tuple[int, Tuple[int, int], int]] = params.benchmark.nBr

    fr: List[Tuple[Objective, int]] = [None] * len(nBr)
    for i, (n, B_range, r) in enumerate(nBr):
        # make sure that all experiments with the same B_range have the same
        # random B vector
        B = n_B_range_to_B_map.setdefault((n, *B_range), compute_B(rng, n, B_range))

        # trim the original bipartite graph G=(V \cup T, E) such that |V| = n
        G = trim_graph(n=n)

        fr[i] = (BudgetAllocation(G=G, B=B, B_range=B_range), r)

    return fr


def get_objective(rng: np.random.Generator,
                  dataset_dir: str,
                  cfg: DictConfig) -> List[Tuple[Objective, int]]:
    """
    Return an instance of the selected set-submodular objective
    :param rng: numpy random generator instance
    :param dataset_dir: datasets main directory
    :param cfg: Hydra configuration dictionary
    """
    objective_name = cfg.obj.name

    print(f'Loading f: {objective_name}\n')
    return OBJ_MAP[objective_name](rng=rng,
                                   params=cfg.obj,
                                   dataset_dir=dataset_dir)

from typing import Iterator, Tuple, Dict, List
import numpy as np
from nptyping import NDArray
from omegaconf import DictConfig
from ..objective import Objective, DemoMonotone, DemoMonotoneSkewed, \
                       DemoNonMonotone, FacilityLocation, BudgetAllocation
from .. import dataset_utils


def compute_B(rng: np.random.Generator,
              n: int, B_range: Tuple[int, int]) -> NDArray[int]:
    """
    Compute the upper-bound array B.
    :param rng: numpy random generator instance
    :param n: size of the ground set
    :param B_range: inclusive range for the values of each entry of B
    """
    low, high = B_range

    return rng.integers(low=low, high=high, size=(n, ), endpoint=True)


OBJ_MAP = {
    'demo_monotone': lambda **kwargs: load_demo_monotone(**kwargs),
    'demo_monotone_skewed': lambda **kwargs: load_demo_monotone_skewed(**kwargs),
    'demo_non_monotone': lambda **kwargs: load_demo_non_monotone(**kwargs),
    'facility_location': lambda **kwargs: load_facility_location(**kwargs),
    'budget_allocation': lambda **kwargs: load_budget_allocation(**kwargs),
}


def load_demo_monotone(rng: np.random.Generator,
                       params,
                       **kwargs) -> List[Tuple[Objective, int]]:
    """
    Generate a random modular, monotone function
    :param rng: numpy random generator instance
    :param params: 'params.demo_monotone' dictionary entry in conf/config.yaml
    """
    nBr: List[Tuple[int, Tuple[int, int], int]] = params.benchmark.nBr

    fr: List[Tuple[Objective, int]] = [None] * len(nBr)
    for i, (n, B_range, r) in enumerate(nBr):
        B = compute_B(rng, n, B_range)
        fr[i] = (DemoMonotone(rng, n=n, B=B, B_range=B_range), r)

    return fr


def load_demo_monotone_skewed(rng: np.random.Generator,
                              params,
                              **kwargs) -> List[Tuple[Objective, int]]:
    """
    Generate a random skewed modular, monotone function
    :param rng: numpy random generator instance
    :param params: 'params.demo_monotone' dictionary entry in conf/config.yaml
    """
    nBr: List[Tuple[int, Tuple[int, int], int]] = params.benchmark.nBr

    fr: List[Tuple[Objective, int]] = [None] * len(nBr)
    for i, (n, B_range, r) in enumerate(nBr):
        B = compute_B(rng, n, B_range)
        fr[i] = (DemoMonotoneSkewed(rng, n=n, B=B, B_range=B_range), r)

    return fr


def load_demo_non_monotone(rng: np.random.Generator,
                           params,
                           **kwargs) -> List[Tuple[Objective, int]]:
    """
    Generate a random modular, non_monotone function
    :param rng: numpy random generator instance
    :param params: 'params.demo_non_monotone' dictionary entry in conf/config.yaml
    """
    nBr: List[Tuple[int, Tuple[int, int], int]] = params.benchmark.nBr

    fr: List[Tuple[Objective, int]] = [None] * len(nBr)
    for i, (n, B_range, r) in enumerate(nBr):
        B = compute_B(rng, n, B_range)
        fr[i] = (DemoNonMonotone(rng, n=n, B=B, B_range=B_range), r)

    return fr


def load_facility_location(rng: np.random.Generator,
                           dataset_dir: str,
                           params,
                           **kwargs) -> List[Tuple[Objective, int]]:
    """
    Generate a random integer-lattice submodular, monotone function that models
    the Facility Location problem.
    :param rng: numpy random generator instance
    :param dataset_dir: datasets main directory
    :param params: 'params.demo_facility_location' dictionary entry in conf/config.yaml
    """
    print(f'Loading Movielens 100k...')
    G = dataset_utils.import_movielens_100k(dataset_dir)
    print(f'...Movielens 100k successfully loaded')
    br: List[Tuple[int, int]] = params.benchmark.br
    
    fr: List[Tuple[Objective, int]] = [None] * len(br)
    for i, (b, r) in enumerate(br):
        fr[i] = (FacilityLocation(G=G, b=b), r)

    return fr


def load_budget_allocation(rng: np.random.Generator,
                           dataset_dir: str,
                           params,
                           **kwargs) -> Iterator[Tuple[Objective, int]]:
    """
    Generate a random integer-lattice DR-submodular, monotone function that models
    the Budget Allocation problem.
    :param rng: numpy random generator instance
    :param dataset_dir: datasets main directory
    :param params: 'params.demo_facility_location' dictionary entry in conf/config.yaml
    """
    # keep track of the vector B generated by the corresponding (n, b_low, b_high)
    n_B_range_to_B_map: Dict[Tuple[int, int, int], NDArray[int]] = dict()

    trim_graph = dataset_utils.import_wikilens_ratings(rng, dataset_dir)
    nBr: List[Tuple[int, Tuple[int, int], int]] = params.benchmark.nBr

    fr: List[Tuple[Objective, int]] = [None] * len(nBr)
    for i, (n, B_range, r) in enumerate(nBr):
        # make sure that all experiments with the same B_range have the same
        # random B vector
        B = n_B_range_to_B_map.setdefault((n, *B_range), compute_B(rng, n, B_range))

        # trim the original bipartite graph G=(V \cup T, E) such that |V| = n
        G = trim_graph(n=n)

        fr[i] = (BudgetAllocation(G=G, B=B, B_range=B_range), r)

    return fr


def get_objective(rng: np.random.Generator,
                  dataset_dir: str,
                  cfg: DictConfig) -> List[Tuple[Objective, int]]:
    """
    Return an instance of the selected set-submodular objective
    :param rng: numpy random generator instance
    :param dataset_dir: datasets main directory
    :param cfg: Hydra configuration dictionary
    """
    objective_name = cfg.obj.name

    print(f'Loading f: {objective_name}\n')
    return OBJ_MAP[objective_name](rng=rng,
                                   params=cfg.obj,
                                   dataset_dir=dataset_dir)

Compute the upper-bound array B.     :param rng: numpy random generator instance     :param n: size of the ground set     :param B_range: inclusive range for the values of each entry of B Generate a random modular, monotone function     :param rng: numpy random generator instance     :param params: 'params.demo_monotone' dictionary entry in conf/config.yaml Generate a random skewed modular, monotone function     :param rng: numpy random generator instance     :param params: 'params.demo_monotone' dictionary entry in conf/config.yaml Generate a random modular, non_monotone function     :param rng: numpy random generator instance     :param params: 'params.demo_non_monotone' dictionary entry in conf/config.yaml Generate a random integer-lattice submodular, monotone function that models     the Facility Location problem.     :param rng: numpy random generator instance     :param dataset_dir: datasets main directory     :param params: 'params.demo_facility_location' dictionary entry in conf/config.yaml Generate a random integer-lattice DR-submodular, monotone function that models     the Budget Allocation problem.     :param rng: numpy random generator instance     :param dataset_dir: datasets main directory     :param params: 'params.demo_facility_location' dictionary entry in conf/config.yaml # keep track of the vector B generated by the corresponding (n, b_low, b_high) # make sure that all experiments with the same B_range have the same # random B vector # trim the original bipartite graph G=(V \cup T, E) such that |V| = n Return an instance of the selected set-submodular objective     :param rng: numpy random generator instance     :param dataset_dir: datasets main directory     :param cfg: Hydra configuration dictionary