from braket.circuits import Circuit, circuit

@circuit.subroutine(register=True)
def simons_oracle(secret_s: str):
    """
    Quantum circuit implementing a particular oracle for Simon's problem. Details of this implementation are
    explained in the Simons Algorithm demo notebook.

    Args:
        secret_s (str): secret string we wish to find
    """
    # Find the index of the first 1 in s, to be used as the flag bit
    flag_bit=secret_s.find('1')
    
    n=len(secret_s)
    
    circ = Circuit()
    # First copy the first n qubits, so that |x>|0> -> |x>|x>
    for i in range(n):
        circ.cnot(i, i+n)
    
    # If flag_bit=-1, s is the all-zeros string, and we do nothing else.
    if flag_bit != -1:
        # Now apply the XOR with s whenever the flag bit is 1.
        for index,bit_value in enumerate(secret_s):
            
            if bit_value not in ['0','1']:
                raise Exception ('Incorrect char \'' + bit_value + '\' in secret string s:' + secret_s)
                
            # XOR with s whenever the flag bit is 1.
            # In terms of gates, XOR means we apply an X gate only whenever the corresponding bit in s is 1.
            # Applying this X only when the flag qubit is 1 means this is a CNOT gate.
            if(bit_value == '1'):
                circ.cnot(flag_bit,index+n)
    return circ