1# Copyright 2021 The Cirq Developers
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15from cirq_web import widget
16
17from cirq.qis.states import bloch_vector_from_state_vector
18from cirq.qis.states import STATE_VECTOR_LIKE
19
20
21class BlochSphere(widget.Widget):
22    def __init__(
23        self,
24        sphere_radius: int = 5,
25        state_vector: STATE_VECTOR_LIKE = None,
26    ):
27        """Initializes a BlochSphere.
28
29        Also initializes it's parent class Widget with the bundle file provided.
30
31        Args:
32            sphere_radius: the radius of the bloch sphere in the three.js diagram.
33            The default value is 5.
34
35            state_vector: a state vector to pass in to be represented.
36        """
37        super().__init__()
38        if sphere_radius <= 0:
39            raise ValueError('You must input a positive radius for the sphere')
40        self.sphere_radius = sphere_radius
41
42        if state_vector is None:
43            raise ValueError('No state vector given in BlochSphere initialization')
44        self.bloch_vector = bloch_vector_from_state_vector(state_vector, 0)
45
46    def get_client_code(self) -> str:
47        return f"""
48        <script>
49        renderBlochSphere('{self.id}', {self.sphere_radius})
50            .addVector({self.bloch_vector[0]}, {self.bloch_vector[1]}, {self.bloch_vector[2]});
51        </script>
52        """
53
54    def get_widget_bundle_name(self) -> str:
55        return 'bloch_sphere.bundle.js'
56