# 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.
"""An MPS simulator.

This is based on this paper:
https://arxiv.org/abs/2002.07730
"""

import dataclasses
import math
from typing import Any, Dict, List, Optional, Sequence, Set, TYPE_CHECKING, Union

import numpy as np
import quimb.tensor as qtn

from cirq import devices, ops, protocols, value
from cirq.sim import simulator_base
from cirq.sim.act_on_args import ActOnArgs

if TYPE_CHECKING:
    import cirq


@dataclasses.dataclass(frozen=True)
class MPSOptions:
    # Some of these parameters are fed directly to Quimb so refer to the documentation for detail:
    # https://quimb.readthedocs.io/en/latest/_autosummary/ \
    #       quimb.tensor.tensor_core.html#quimb.tensor.tensor_core.tensor_split

    # How to split the tensor. Refer to the Quimb documentation for the exact meaning.
    method: str = 'svds'
    # If integer, the maxmimum number of singular values to keep, regardless of ``cutoff``.
    max_bond: Optional[int] = None
    # Method with which to apply the cutoff threshold. Refer to the Quimb documentation.
    cutoff_mode: str = 'rsum2'
    # The threshold below which to discard singular values. Refer to the Quimb documentation.
    cutoff: float = 1e-6
    # Because the computation is approximate, the sum of the probabilities is not 1.0. This
    # parameter is the absolute deviation from 1.0 that is allowed.
    sum_prob_atol: float = 1e-3


class MPSSimulator(
    simulator_base.SimulatorBase[
        'MPSSimulatorStepResult', 'MPSTrialResult', 'MPSState', 'MPSState'
    ],
):
    """An efficient simulator for MPS circuits."""

    def __init__(
        self,
        noise: 'cirq.NOISE_MODEL_LIKE' = None,
        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
        simulation_options: MPSOptions = MPSOptions(),
        grouping: Optional[Dict['cirq.Qid', int]] = None,
    ):
        """Creates instance of `MPSSimulator`.

        Args:
            noise: A noise model to apply while simulating.
            seed: The random seed to use for this simulator.
            simulation_options: Numerical options for the simulation.
            grouping: How to group qubits together, if None all are individual.

        Raises:
            ValueError: If the noise model is not unitary or a mixture.
        """
        self.init = True
        noise_model = devices.NoiseModel.from_noise_model_like(noise)
        if not protocols.has_mixture(noise_model):
            raise ValueError(f'noise must be unitary or mixture but was {noise_model}')
        self.simulation_options = simulation_options
        self.grouping = grouping
        super().__init__(
            noise=noise,
            seed=seed,
        )

    def _create_partial_act_on_args(
        self,
        initial_state: Union[int, 'MPSState'],
        qubits: Sequence['cirq.Qid'],
        classical_data: 'cirq.ClassicalDataStore',
    ) -> 'MPSState':
        """Creates MPSState args for simulating the Circuit.

        Args:
            initial_state: The initial state for the simulation in the
                computational basis. Represented as a big endian int.
            qubits: Determines the canonical ordering of the qubits. This
                is often used in specifying the initial state, i.e. the
                ordering of the computational basis states.
            classical_data: The shared classical data container for this
                simulation.

        Returns:
            MPSState args for simulating the Circuit.
        """
        if isinstance(initial_state, MPSState):
            return initial_state

        return MPSState(
            qubits=qubits,
            prng=self._prng,
            simulation_options=self.simulation_options,
            grouping=self.grouping,
            initial_state=initial_state,
            classical_data=classical_data,
        )

    def _create_step_result(
        self,
        sim_state: 'cirq.OperationTarget[MPSState]',
    ):
        return MPSSimulatorStepResult(sim_state)

    def _create_simulator_trial_result(
        self,
        params: 'cirq.ParamResolver',