From f3b0ed981e8e95e956001bfad06ac63878dc570e Mon Sep 17 00:00:00 2001
From: Lauren Capelluto <laurecap@amazon.com>
Date: Wed, 1 Nov 2023 17:22:45 -0400
Subject: [PATCH] Clean up

---
 .../experimental/autoqasm/program/program.py  | 28 +++++++++++++------
 .../experimental/autoqasm/test_parameters.py  |  3 ++
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/src/braket/experimental/autoqasm/program/program.py b/src/braket/experimental/autoqasm/program/program.py
index 77b5a2a42..7c213c97f 100644
--- a/src/braket/experimental/autoqasm/program/program.py
+++ b/src/braket/experimental/autoqasm/program/program.py
@@ -16,12 +16,14 @@
 
 import contextlib
 import threading
+import copy
 from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Optional, Union
 
 import oqpy.base
+from openqasm3 import ast
 
 from braket.circuits.free_parameter import FreeParameter
 from braket.circuits.free_parameter_expression import FreeParameterExpression
@@ -123,12 +125,17 @@ def with_calibrations(self, gate_calibrations: Union[Callable, list[Callable]])
         combined_oqpy_program += self._oqpy_program
         return Program(combined_oqpy_program, has_pulse_control=True)
 
-    def make_bound_program(self, param_values: dict[str, float]) -> Program:  # Number
-        # FIXME
-        import copy
-        from openqasm3 import ast
+    def make_bound_program(self, param_values: dict[str, float]) -> Program:
+        """TODO
 
+        Args:
+            param_values (dict[str, float]): TODO
+        Returns:
+            TODO
+        """
+        # We have to copy the program so that we don't modify the original, unbound program
         oqpy_program_copy = copy.deepcopy(self._oqpy_program)
+        # Let's break early if we've processed all the parameter assignments
         params_to_process = set(param_values.keys())
         for state in oqpy_program_copy.stack:
             for i in range(len(state.body)):
@@ -137,13 +144,16 @@ def make_bound_program(self, param_values: dict[str, float]) -> Program:  # Numb
                     name = inst.identifier.name
                     target = oqpy_program_copy.declared_vars[name]
                     target.init_expression = param_values[name]
+                    # TODO: how do I create the right statement directly?
                     inst = oqpy.Program().declare(target).stack[0].body[0]
                     state.body[i] = inst
                     params_to_process.remove(name)
-                if params_to_process == set():
-                    break
+                    if params_to_process == set():
+                        break
             else:
                 continue
+            # If we broke out of the inner loop, we should break out of the outer loop, too,
+            # because there's no more work to do
             break
 
         return Program(oqpy_program_copy, self._has_pulse_control)
@@ -332,21 +342,21 @@ def register_parameter(self, name: str) -> None:
             name (str): The identifier for the parameter.
         """
         if name not in self._free_parameters:
-            self._free_parameters[name] = (oqpy.FloatVar("input", name=name), FreeParameter(name))
+            self._free_parameters[name] = oqpy.FloatVar("input", name=name)
 
     def get_free_parameters(self) -> list[oqpy.FloatVar]:
         """Return a list of named oqpy.Vars that are used as free parameters in the program."""
-        return list(val[0] for val in self._free_parameters.values())
+        return list(self._free_parameters.values())
 
     def add_io_declarations(self) -> None:
         free_parameters = self.get_free_parameters()
         root_oqpy_program = self.get_oqpy_program(scope=ProgramScope.MAIN)
         for parameter in free_parameters[::-1]:
+            # TODO: is it possible to save a pointer to the declaration statement?
             root_oqpy_program.declare(
                 parameter,
                 to_beginning=True,
             )
-            # TODO: is it possible to save a pointer to the declaration?
 
     def get_target_device(self) -> Optional[Device]:
         """Return the target device for the program, as specified by the user.
diff --git a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
index 2b9844545..00bcdbb37 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_parameters.py
@@ -391,3 +391,6 @@ def rx_alpha(int[32] qubit, float[64] theta) {
 rx_alpha(2, alpha);
 rx_alpha(2, beta);"""
     assert bound_prog.to_ir() == expected
+
+
+# TODO: test with pulse and gates