##
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
##
import warnings
import logging

from typing import TYPE_CHECKING, Union, Any, Optional
from enum import Enum
from azure.quantum.job.base_job import ContentType
from azure.quantum import Workspace, Job
from azure.quantum.job.base_job import DEFAULT_TIMEOUT
from azure.quantum.target.target import Target

if TYPE_CHECKING:
    from azure.quantum.optimization.problem import Problem

logger = logging.getLogger(__name__)

__all__ = [
    "RangeSchedule",
    "Solver",
]



class RangeSchedule:
    def __init__(self, schedule_type: str, initial: float, final: float):
        """Constructor of RangeSchedule for solvers.

        :param schedule_type:
            str, type of the RangeSchedule
            currently only supports 'linear' and 'geometric'.
        :param initial:
            float, initial value of RangeSchedule
        :param final:
            float, final value of the RangeSchedule
        """
        self.schedule_type = schedule_type
        self.initial = initial
        self.final = final
        if not (
            self.schedule_type == "linear" or self.schedule_type == "geometric"
        ):
            raise ValueError(
                '"schedule_type" must be "linear" or "geometric"!'
            )

class Solver(Target):
    def __init__(
        self,
        workspace: Workspace,
        name: str,
        provider_id: str,
        input_data_format="microsoft.qio.v2",
        output_data_format="microsoft.qio-results.v2",
        nested_params: bool = True,
        force_str_params: bool = False,
        params: dict = None,
        content_type : Optional[ContentType] = ContentType.json,
        **kwargs
    ):
        self.provider_id = provider_id
        self.nested_params = nested_params
        self.force_str_params = force_str_params
        self.params = {"params": {}} if nested_params else {}
        params = params or {}
        name = self._switch_name(name, **params)

        super().__init__(
            workspace=workspace,
            name=name,
            input_data_format=input_data_format,
            output_data_format=output_data_format,
            provider_id=provider_id,
            content_type = content_type,
            encoding="gzip",
            **kwargs
        )
 
        self._set_params(**params)
        self._check_params(**params)


    """Constants that define thresholds for submission warnings
    """
    SWEEPS_WARNING = 10000
    TIMEOUT_WARNING = 600

    def _switch_name(self, name, **params):
        return name

    def _check_params(self, **params):
        pass

    @staticmethod
    def _encode_input_data(data: "Problem") -> bytes:
        return data.to_blob()

    def submit(
        self, problem: Union[str, "Problem"]) -> Job:
        """Submits a job to execution to the associated
        Azure Quantum Workspace.

        :param problem:
            The Problem to solve. It can be an instance of a Problem,
            or the URL of an Azure Storage Blob where the serialized version
            of a Problem has been uploaded.
        """
        from azure.quantum.optimization import Problem
        if isinstance(problem, Problem):
            return super().submit(
                input_data=problem,
                name=problem.name,
                input_params=self.params,
                blob_name="inputData",
                content_type = problem.content_type
            )

        else:
            if hasattr(problem, "uploaded_blob_uri"):
                name = problem.name
                problem_uri = problem.uploaded_blob_uri

            elif isinstance(problem, str):
                name = "Optimization problem"
                problem_uri = problem
            
            else:
                raise ValueError("Cannot submit problem: should be of type str, Problem or have uploaded_blob_uri attribute.")

            # Create job from storage URI
            job = Job.from_storage_uri(
                workspace=self.workspace,
                name=name,
                target=self.name,
                input_data_uri=problem_uri,
                provider_id=self.provider_id,
                input_data_format=self.input_data_format,
                output_data_format=self.output_data_format,
                input_params=self.params,
                session_id=self.get_latest_session_id()
            )

        return job

    def optimize(self, problem: Union[str, "Problem"], timeout_secs: int=DEFAULT_TIMEOUT):
        """[Submits the Problem to the associated
            Azure Quantum Workspace and get the results.

        :param problem:
            The Problem to solve. It can be an instance of a Problem,
            or the URL of an Azure Storage Blob where the serialized version
            of a Problem has been uploaded.
        :type problem: Union[str, Problem]
        :param timeout_secs: Timeout in seconds, defaults to 300
        :type timeout_secs: int
        :return: Job results
        :rtype: dict
        """
        if not isinstance(problem, str):
            self.check_submission_warnings(problem)

        job = self.submit(problem)
        logger.info(f"Submitted job: '{job.id}'")

        return job.get_results(timeout_secs=timeout_secs)

    def set_one_param(self, name: str, value: Any):
        if value is not None:
            params = (
                self.params["params"] if self.nested_params else self.params
            )
            params[name] = str(value) if self.force_str_params else value

    def _set_params(self, **params):
        for param_name, param in params.items():
            self.set_one_param(param_name, param)

    def set_number_of_solutions(self, number_of_solutions: int):
        """Sets the number of solutions to return.

        Applies to all solvers.
        Default value is 1 if not supplied.

        :param number_of_solutions:
            Number of solutions to return. Must be a positive integer.
        """
        self.set_one_param("number_of_solutions", number_of_solutions)

    def check_submission_warnings(self, problem: "Problem"):
        # print a warning if we suspect the job
        # may take long based on its configured properties.
        is_large_problem = problem.is_large()
        if is_large_problem:
            if self.nested_params and "sweeps" in self.params["params"]:
                sweeps = int(self.params["params"]["sweeps"])
                # if problem is large and sweeps is large, warn.
                if sweeps >= Solver.SWEEPS_WARNING:
                    logger.warning(
                        f"There is a large problem submitted and \
                        a large number of sweeps ({sweeps}) configured. \
                        This submission could result in a long runtime."
                    )

        # do a timeout check if param-free, to warn
        # new users who accidentally set high timeout values.
        if self.nested_params and "timeout" in self.params["params"]:
            timeout = int(self.params["params"]["timeout"])
            if timeout >= Solver.TIMEOUT_WARNING:
                logger.warning(
                    f"A large timeout has been set for this submission \
                        ({timeout}). \
                        If this is intended, disregard this warning. \
                        Otherwise, consider cancelling the job \
                        and resubmitting with a lower timeout."
                )
    
    class ScheduleEvolution(Enum):
        INCREASING = 1
        DECREASING = 2

    def check_set_schedule(
        self,
        schedule_name: str,
        schedule_value: RangeSchedule,
        evolution: Optional[ScheduleEvolution] = None,
        lower_bound_exclusive: Optional[float] = None,
        lower_bound_inclusive: Optional[float] = None,
    ):
        """Set the parameter `schedule_name` from the value `schedule_value`
        and check that it is of type `RangeSchedule` and each end satisfies
        the specified bound.

        :param schedule_name:
            Name of the schedule parameter.
        :param schedule_value:
            Schedule value to be assigned and checked.
        :param evolution
            Expected schedule evolution (INCREASING or DECREASING)
        :lower_bound_exclusive:
            Exclusive lower bound for both ends of the schedule, optional.
        :lower_bound_inclusive:
            Inclusive lower bound for both ends of the schedule, optional.
        """
        if not (schedule_value is None):
            if not isinstance(schedule_value, RangeSchedule):
                raise ValueError(
                    f"{schedule_name} must be of type RangeSchedule; found"
                    f" type({schedule_name})={type(schedule_value).__name__}."
                )
            schedule_param = {
                "type": schedule_value.schedule_type,
                "initial": schedule_value.initial,
                "final": schedule_value.final,
            }
            if evolution is not None:
                if (evolution == self.ScheduleEvolution.INCREASING and
                        schedule_value.initial > schedule_value.final):
                    raise ValueError(
                            f"Schedule for {schedule_name} must be increasing;"
                            f" found {schedule_name}.initial"
                            f"={schedule_value.initial}"
                            f" > {schedule_value.final}"
                            f"={schedule_name}.final."
                    )
                if (evolution == self.ScheduleEvolution.DECREASING and
                        schedule_value.initial < schedule_value.final):
                    raise ValueError(
                            f"Schedule for {schedule_name} must be decreasing;"
                            f" found {schedule_name}.initial"
                            f"={schedule_value.initial}"
                            f" < {schedule_value.final}"
                            f"={schedule_name}.final."
                    )
            self.check_limit(
                    f"{schedule_name}.initial",
                    schedule_value.initial,
                    lower_bound_exclusive=lower_bound_exclusive,
                    lower_bound_inclusive=lower_bound_inclusive)
            self.check_limit(
                    f"{schedule_name}.final",
                    schedule_value.final,
                    lower_bound_exclusive=lower_bound_exclusive,
                    lower_bound_inclusive=lower_bound_inclusive)
            self.set_one_param(schedule_name, schedule_param)

    def check_set_positive_int(
            self,
            parameter_name: str,
            parameter_value: int):
        """Set the parameter `parameter_name` from the value `parameter_value`
        and check that it is a positive integer.

        :param parameter_name:
            Name of the parameter.
        :param parameter_value:
            Value to be assigned and checked.
        """
        if not (parameter_value is None):
            if not isinstance(parameter_value, int):
                raise ValueError(
                        f"{parameter_name} must be of type int; found"
                        f"type({parameter_name})"
                        f"={type(parameter_value).__name__}.")
            if parameter_value <= 0:
                raise ValueError(
                        f"{parameter_name} must be positive; found "
                        f"{parameter_name}={parameter_value}.")
            self.set_one_param(parameter_name, parameter_value)

    def check_set_float(
        self,
        parameter_name: str,
        parameter_value: float,
        lower_bound_exclusive: Optional[float] = None,
        lower_bound_inclusive: Optional[float] = None,
    ):
        """Set the parameter `parameter_name` from the value `parameter_value`
        and check that it has a float value satisfying bounds.

        :param parameter_name:
            Name of the parameter.
        :param parameter_value:
            Value to be assigned and checked.
        :lower_bound_exclusive:
            Exclusive lower bound to check parameter_value against, optional.
        :lower_bound_inclusive:
            Inclusive lower bound to check parameter_value against, optional.
        """
        if not (parameter_value is None):
            if not (isinstance(parameter_value, float) or
                    isinstance(parameter_value, int)):
                raise ValueError(f"{parameter_name} must be a float!")
            else:
                self.check_limit(
                        parameter_name=parameter_name,
                        parameter_value=parameter_value,
                        lower_bound_exclusive=lower_bound_exclusive,
                        lower_bound_inclusive=lower_bound_inclusive)
                self.set_one_param(parameter_name, parameter_value)

    def check_limit(
        self,
        parameter_name: str,
        parameter_value: Optional[float],
        lower_bound_exclusive: Optional[float] = None,
        lower_bound_inclusive: Optional[float] = None,
    ):
        """Check whether `parameter_value` satisfies a lower bound.

        :param parameter_name:
            Name of the parameter.
        :param parameter_value:
            Value to be checked.
        :lower_bound_exclusive:
            Exclusive lower bound to check parameter_value against, optional.
        :lower_bound_inclusive:
            Inclusive lower bound to check parameter_value against, optional.
        """
        if not (parameter_value is None):
            if (lower_bound_exclusive is not None and
                    parameter_value <= lower_bound_exclusive):
                raise ValueError(
                    f"{parameter_name} must be greater than "
                    f"{lower_bound_exclusive}; "
                    f"found {parameter_name}={parameter_value}."
                )
            if (lower_bound_inclusive is not None and
                    parameter_value < lower_bound_inclusive):
                raise ValueError(
                    f"{parameter_name} must be greater equal "
                    f"{lower_bound_inclusive}; found "
                    f"{parameter_name}={parameter_value}."
                )

