# Copyright 2019 The Cirq Developers
#
# 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
#
#     https://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.

"""Sampling/simulation methods that delegate to appropriate simulators.

Filename is a reference to multiplexing.
"""

from typing import cast, List, Optional, Sequence, Type, TYPE_CHECKING, Union

import numpy as np

from cirq import circuits, protocols, study, devices, ops, value
from cirq._doc import document
from cirq.sim import sparse_simulator, density_matrix_simulator
from cirq.sim.clifford import clifford_simulator
from cirq.transformers import measurement_transformers

if TYPE_CHECKING:
    import cirq

CIRCUIT_LIKE = Union[circuits.Circuit, ops.Gate, ops.OP_TREE]
document(
    CIRCUIT_LIKE,  # type: ignore
    """A `circuits.Circuit` or a value that can be trivially converted into it:
        a gate, an operation, and a list or tree of operations.
    """,
)


def _is_clifford_circuit(program: 'cirq.Circuit') -> bool:
    return all(
        clifford_simulator.CliffordSimulator.is_supported_operation(op)
        for op in program.all_operations()
    )


def sample(
    program: 'cirq.Circuit',
    *,
    noise: 'cirq.NOISE_MODEL_LIKE' = None,
    param_resolver: Optional['cirq.ParamResolver'] = None,
    repetitions: int = 1,
    dtype: Type[np.number] = np.complex64,
    seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
) -> 'cirq.Result':
    """Simulates sampling from the given circuit.

    Args:
        program: The circuit to sample from.
        noise: Noise model to use while running the simulation.
        param_resolver: Parameters to run with the program.
        repetitions: The number of samples to take.
        dtype: The `numpy.dtype` used by the simulation. Typically one of
            `numpy.complex64` or `numpy.complex128`.
            Favors speed over precision by default, i.e. uses `numpy.complex64`.
        seed: The random seed to use for this simulator.
    """
    noise_model = devices.NoiseModel.from_noise_model_like(noise)

    # State vector simulation is much faster, but only works if no randomness.
    if noise_model == devices.NO_NOISE:
        if _is_clifford_circuit(program):
            # If all non-measurement operations are clifford, use the Clifford
            # simulator.
            return clifford_simulator.CliffordSimulator(seed=seed).run(
                program, param_resolver=param_resolver, repetitions=repetitions
            )
        if protocols.has_unitary(program):