# Copyright 2019 IBM Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
import itertools
import warnings
import random
import math
from collections import ChainMap

from lale.util.Visitor import Visitor, accept
from lale.search.search_space import SearchSpace, SearchSpaceObject, SearchSpaceConstant, SearchSpaceEnum, SearchSpaceSum, SearchSpaceProduct, SearchSpacePrimitive, SearchSpaceArray, SearchSpaceOperator, should_print_search_space
from lale.search.schema2search_space import op_to_search_space
from lale.search.PGO import PGO
from lale.sklearn_compat import nest_all_HPparams, nest_choice_all_HPparams, DUMMY_SEARCH_SPACE_GRID_PARAM_NAME, discriminant_name, make_indexed_name, make_array_index_name, structure_type_name, structure_type_list, structure_type_tuple, structure_type_dict

from lale.operators import PlannedOperator, OperatorChoice, PlannedIndividualOp, PlannedPipeline, Operator

SearchSpaceGrid = Dict[str,SearchSpacePrimitive]

def search_space_grid_to_string(grid:SearchSpaceGrid)->str:
    return "{" + ";".join(f"{k}->{str(v)}" for k,v in grid.items()) + "}"

def search_space_grids_to_string(grids:List[SearchSpaceGrid])->str:
    return "|".join(search_space_grid_to_string(grid) for grid in grids)

def get_search_space_grids( op:'PlannedOperator', 
                            num_grids:Optional[float]=None, 
                            pgo:Optional[PGO]=None)->List[SearchSpaceGrid]:
    """ Top level function: given a lale operator, returns a list of hp grids.
    Parameters
    ----------
    op : The lale PlannedOperator
    num_grids: integer or float, optional
        if set to an integer => 1, it will determine how many parameter grids will be returned (at most)
        if set to an float between 0 and 1, it will determine what fraction should be returned
        note that setting it to 1 is treated as in integer.  To return all results, use None
    """
    all_parameters = op_to_search_space_grids(op, pgo=pgo)
    if should_print_search_space("true", "all", "search_space_grids", "grids"):
        name = op.name()
        if not name:
            name = "an operator"
        print(f"search space grids for {name}:\n{search_space_grids_to_string(all_parameters)}")
    if num_grids is None:
        return all_parameters
    else:
        if num_grids <= 0:
            warnings.warn(f"get_search_space_grids(num_grids={num_grids}) called with a non-positive value for lale_num_grids")
            return []
        if num_grids >= 1:
            samples = math.ceil(num_grids)
            if samples >= len(all_parameters):
                return all_parameters
            else:
                warnings.warn(f"get_search_space_grids(num_grids={num_grids}) sampling {math.ceil(num_grids)}/{len(all_parameters)}")
                return random.sample(all_parameters, math.ceil(num_grids))
        else:
            samples = round(len(all_parameters)*num_grids)
            warnings.warn(f"get_search_space_grids(num_grids={num_grids}) sampling {samples}/{len(all_parameters)}")
            return random.sample(all_parameters, samples)

def search_space_to_grids(hp:SearchSpace)->List[SearchSpaceGrid]:
    return SearchSpaceToGridVisitor.run(hp)

def op_to_search_space_grids(op:PlannedOperator,
                             pgo:Optional[PGO]=None)->List[SearchSpaceGrid]:
    search_space = op_to_search_space(op, pgo=pgo)
    grids = search_space_to_grids(search_space)
    return grids

# lets handle the general case
SearchSpaceGridInternalType = Union[List[SearchSpaceGrid], SearchSpacePrimitive]

class SearchSpaceToGridVisitor(Visitor):
    @classmethod
    def run(cls, space:SearchSpace)->List[SearchSpaceGrid]:
        visitor = cls()
        space_:Any = space
        grids:SearchSpaceGridInternalType = space_.accept(visitor)
        fixed_grids = cls.fixupDegenerateSearchSpaces(grids)
        return fixed_grids

    @classmethod
    def fixupDegenerateSearchSpaces(cls, space:SearchSpaceGridInternalType)->List[SearchSpaceGrid]:
        if isinstance(space, SearchSpacePrimitive):
            return [{DUMMY_SEARCH_SPACE_GRID_PARAM_NAME:space}]
        else:
            return space

    def __init__(self):
        super(SearchSpaceToGridVisitor, self).__init__()

    def visitSearchSpacePrimitive(self, space:SearchSpacePrimitive)->SearchSpacePrimitive:
        return space
    
    visitSearchSpaceEnum = visitSearchSpacePrimitive

    visitSearchSpaceConstant = visitSearchSpaceEnum
    visitSearchSpaceBool = visitSearchSpaceEnum

    visitSearchSpaceNumber = visitSearchSpacePrimitive

    def _searchSpaceList(self, space:SearchSpaceArray, *, size:int)->List[SearchSpaceGrid]:
        sub_spaces = space.items(max=size)

        param_grids:List[List[SearchSpaceGrid]] = \
            [nest_all_HPparams(str(index), self.fixupDegenerateSearchSpaces(accept(sub,self))) for index,sub in enumerate(sub_spaces)]

        param_grids_product:Iterable[Iterable[SearchSpaceGrid]] = itertools.product(*param_grids)
        chained_grids:List[SearchSpaceGrid] = [
            dict(ChainMap(*gridline,)) for gridline in param_grids_product]

        if space.is_tuple:
            st_val = structure_type_tuple
        else:
            st_val = structure_type_list

        discriminated_grids:List[SearchSpaceGrid]=[{**d, structure_type_name:SearchSpaceConstant(st_val)} for d in chained_grids]


        return discriminated_grids

    def visitSearchSpaceArray(self, space:SearchSpaceArray)->List[SearchSpaceGrid]:
        if space.minimum == space.maximum:
            return self._searchSpaceList(space, size=space.minimum)
        else:
            ret:List[SearchSpaceGrid] = []
            for i in range(space.minimum, space.maximum+1):
                ret.extend(self._searchSpaceList(space, size=i))
            return ret

    def visitSearchSpaceObject(self, space:SearchSpaceObject)->List[SearchSpaceGrid]:
        keys = space.keys
        keys_len = len(keys)
        final_choices:List[SearchSpaceGrid] = []
        for c in space.choices:
            assert keys_len == len(c)
            kvs_complex:List[List[SearchSpaceGrid]] = []
            kvs_simple:SearchSpaceGrid = {}
            for k,v in zip(keys, c):
                vspace:Union[List[SearchSpaceGrid],SearchSpacePrimitive] = v.accept(self)
                if isinstance(vspace, SearchSpacePrimitive):
                    kvs_simple[k] = vspace
                else:
                    nested_vspace:List[SearchSpaceGrid] = nest_all_HPparams(k, vspace)
                    if nested_vspace:
                        kvs_complex.append(nested_vspace)
            nested_space_choices:Iterable[Iterable[SearchSpaceGrid]] = itertools.product(*kvs_complex)
            nested_space_choices_lists:List[List[SearchSpaceGrid]] = list(map((lambda x: list(x)),nested_space_choices))
            nested_space_choices_filtered:List[List[SearchSpaceGrid]] = [
                l for l in nested_space_choices_lists if l]
            if nested_space_choices_filtered:
                chained_grids:Iterable[SearchSpaceGrid] = [
                  dict(ChainMap(*nested_choice, kvs_simple)) for nested_choice in nested_space_choices_filtered]
                final_choices.extend(chained_grids)
            else:
                final_choices.append(kvs_simple)
        return final_choices

    def visitSearchSpaceSum(self, op:SearchSpaceSum)->SearchSpaceGridInternalType:
        sub_spaces:List[SearchSpace] = op.sub_spaces

        sub_grids:Iterable[SearchSpaceGridInternalType] = [
            accept(cur_space, self) for cur_space in sub_spaces]

        if len(sub_spaces) == 1:
            return list(sub_grids)[0]
        else:
            fixed_grids:Iterable[List[SearchSpaceGrid]] = (SearchSpaceToGridVisitor.fixupDegenerateSearchSpaces(grid) for grid in sub_grids)
            final_grids:List[SearchSpaceGrid] = []
            for i, grids in enumerate(fixed_grids):
                if not grids:
                    grids = [{}]
                else:
                    # we need to add in this nesting
                    # in case a higher order operator directly contains
                    # another
                    grids = nest_choice_all_HPparams(grids)

                discriminated_grids:List[SearchSpaceGrid]=[{**d, discriminant_name:SearchSpaceConstant(i)} for d in grids]
                final_grids.extend(discriminated_grids)
            return final_grids

    def visitSearchSpaceProduct(self, op:SearchSpaceProduct)->SearchSpaceGridInternalType:

        sub_spaces = op.get_indexed_spaces()

        param_grids:List[List[SearchSpaceGrid]] = [nest_all_HPparams(make_indexed_name(name, index), self.fixupDegenerateSearchSpaces(accept(space,self))) for name,index,space in sub_spaces]

        param_grids_product:Iterable[Iterable[SearchSpaceGrid]] = itertools.product(*param_grids)
        chained_grids:List[SearchSpaceGrid] = [
            dict(ChainMap(*gridline)) for gridline in param_grids_product]

        return chained_grids

    def visitSearchSpaceOperator(self, op:SearchSpaceOperator)->SearchSpaceGridInternalType:
        return accept(op.sub_space, self)
