import copy

from typing import Dict, List, Union

from clipped.formatting import Printer

from polyaxon.exceptions import PolyaxonSchemaError
from polyaxon.polyaxonfile.specs.libs.parser import PolyaxonfileParser
from polyaxon.polyflow import V1CompiledOperation, V1MatrixKind, V1Operation, V1Param


def get_op_from_schedule(
    content: str,
    compiled_operation: V1CompiledOperation,
) -> V1Operation:
    op_spec = V1Operation.read(content)
    op_spec.conditions = None
    op_spec.schedule = None
    op_spec.events = None
    op_spec.dependencies = None
    op_spec.trigger = None
    op_spec.build = None
    op_spec.skip_on_upstream_skip = None
    op_spec.cache = compiled_operation.cache
    op_spec.queue = compiled_operation.queue
    op_spec.component.inputs = compiled_operation.inputs
    op_spec.component.outputs = compiled_operation.outputs
    op_spec.component.run = compiled_operation.run
    return op_spec


def get_ops_from_suggestions(
    content: str,
    compiled_operation: V1CompiledOperation,
    suggestions: List[Dict],
) -> List[V1Operation]:
    io_keys = compiled_operation.get_io_names()

    def has_param(k: str):
        if not compiled_operation.matrix:
            return None
        return k not in io_keys

    op_content = V1Operation.read(content)  # TODO: Use construct
    for suggestion in suggestions:
        params = {
            k: V1Param(
                value=PolyaxonfileParser.parse_expression(v, {}),
                context_only=has_param(k),
            )
            for (k, v) in suggestion.items()
        }
        op_spec = copy.deepcopy(op_content)
        op_spec.matrix = None
        op_spec.conditions = None
        op_spec.schedule = None
        op_spec.events = None
        op_spec.dependencies = None
        op_spec.trigger = None
        op_spec.build = None
        op_spec.is_approved = None
        op_spec.skip_on_upstream_skip = None
        op_spec.cache = compiled_operation.cache
        op_spec.queue = compiled_operation.queue
        op_spec.params = params
        op_spec.component.inputs = compiled_operation.inputs
        op_spec.component.outputs = compiled_operation.outputs
        op_spec.component.run = compiled_operation.run
        yield op_spec


def get_eager_matrix_operations(
    content: str,
    compiled_operation: V1CompiledOperation,
    is_cli: bool = False,
) -> List[V1Operation]:
    is_supported_in_eager_mode(compiled_operation)

    try:
        import numpy as np
    except ImportError as e:
        if is_cli:
            Printer.error(
                "numpy is required for the eager mode!",
                sys_exit=True,
            )
        raise e

    from hypertune.search_managers.grid_search.manager import GridSearchManager
    from hypertune.search_managers.mapping.manager import MappingManager
    from hypertune.search_managers.random_search.manager import RandomSearchManager

    if compiled_operation.has_random_search_matrix:
        suggestions = RandomSearchManager(compiled_operation.matrix).get_suggestions()
    elif compiled_operation.has_grid_search_matrix:
        suggestions = GridSearchManager(compiled_operation.matrix).get_suggestions()
    elif compiled_operation.has_mapping_matrix:
        suggestions = MappingManager(compiled_operation.matrix).get_suggestions()
    else:
        raise PolyaxonSchemaError(
            "Received a bad configuration, eager mode not supported, "
            "I should not be here!"
        )
    if is_cli:
        Printer.heading("Creating {} operations".format(len(suggestions)))
    return get_ops_from_suggestions(
        content=content, compiled_operation=compiled_operation, suggestions=suggestions
    )


def is_supported_in_eager_mode(spec: Union[V1Operation, V1CompiledOperation]):
    if not spec.matrix:
        if spec.component and spec.component.run:
            raise PolyaxonSchemaError(
                "This operation with runtime `{}` "
                "is not supported in eager mode".format(spec.component.run.kind)
            )
        else:
            raise PolyaxonSchemaError(
                "Received a bad configuration, eager mode not supported"
            )

    if spec.get_matrix_kind() not in V1MatrixKind.eager_values():
        raise PolyaxonSchemaError(
            "This operation is defining a matrix kind `{}` "
            "which is not supported in eager mode".format(spec.get_matrix_kind())
        )
