"""
/*********************************************************************
* Copyright (c) 2023 the Qrisp Authors
*
* This program and the accompanying materials are made
* available under the terms of the Eclipse Public License 2.0
* which is available at https://www.eclipse.org/legal/epl-2.0/
*
* SPDX-License-Identifier: EPL-2.0
**********************************************************************/
"""

# -*- coding: utf-8 -*-

import threading
from qrisp.simulator.quantum_state import QuantumState

import numpy as np
thread_limit = 1

#This class describes a set of decoherent quantum states
class ImpureQuantumState:
    
    
    #Initially we start we a single quantum state
    def __init__(self, n, clbit_amount = 0):
        self.states = [QuantumState(n)]
        self.cl_prob = [1]
        self.outcome_list = [np.zeros(clbit_amount, dtype = np.byte)]
        self.collected_measurements = []
        
    
    
    def apply_operation(self, operation, qubits):
        
        #Perform operation application in parallel. This is possible because many
        #functions in here release the GIL (such as numpy's np.dot)
        if thread_limit > 1 and len(self.states) > 1:
            
            threads = []
            
            for i in range(len(self.states)):
                mp_wrapper = lambda : self.states[i].apply_operation(operation, qubits)
                
                thr = threading.Thread(target = mp_wrapper)
                
                if len(threads) >= thread_limit:
                    threads[0].join()
                    threads.pop(0)
                
                thr.start()
                threads.append(thr)
                
            
            [thr.join() for thr in threads]
            
        else:
            
            for i in range(len(self.states)):
                self.states[i].apply_operation(operation, qubits)
            

    def conditionally_apply_operation(self, operation, qubits, clbit):
        
        #Perform operation application in parallel. This is possible because many
        #functions in here release the GIL (such as numpy's np.dot)
        
        if thread_limit > 1:
            
            threads = []
            
            for i in range(len(self.states)):
                
                if self.outcome_list[clbit] == "0":
                    continue
                
                mp_wrapper = lambda : self.states[i].apply_operation(operation, qubits)
                
                if len(threads) >= thread_limit:
                    threads[0].join()
                    threads.pop(0)
                
                
                thr = threading.Thread(target = mp_wrapper)
                thr.start()
                
                threads.append(thr)
            
            [thr.join() for thr in threads]
            
        else:
            
            for i in range(len(self.states)):
                if self.outcome_list[clbit] == "0":
                    continue

                
                self.states[i].apply_operation(operation, qubits)

        
    #Perform reset operation
    #We perform the reset operation by measuring but not tracking the measurement result
    def reset(self, qubit_index, keep_outcome = False):
        
        if thread_limit > 1:
            
            threads = []
            
            for i in range(len(self.states)):
                mp_wrapper = lambda : self.states[i].measure(qubit_index, keep_outcome)
                thr = threading.Thread(target = mp_wrapper)
                
                if len(threads) > thread_limit:
                    threads[0].join()
                    threads.pop(0)
                
                
                thr.start()
                threads.append(thr)
                
            
        else:
            
            for i in range(len(self.states)):
                self.states[i].measure(qubit_index, keep_outcome)
        
        
        #Generate new list of states and the corresponding classical propabilities + the outcome bitstrings
        new_states = []
        new_cl_prob = []
        new_outcome_list = []
        
        for i in range(len(self.states)):
            
            #Wait for measurement threads to finish
            if thread_limit > 1:
                if len(threads):
                    threads[0].join()
                    threads.pop(0)

            #Get measurement outcome
            #p_0 represents the probability to measure a 0,
            #state_0 represents the state after measuring a 0
            p_0, state_0, p_1, state_1 = self.states[i].last_mes_outcome
            #Log results
            if p_0 != 0:
                
                #Add the new decoherent state to the list of states
                new_states.append(state_0)
                
                #Update its classical probability
                new_cl_prob.append(p_0*self.cl_prob[i])
                
                #The bitstring stays the same
                new_outcome = self.outcome_list[i].copy()
                new_outcome_list.append(new_outcome)
                
            #Perform the same thing for the outcome 1
            if p_1 != 0:
                new_states.append(state_1)
                new_cl_prob.append(p_1*self.cl_prob[i])
                new_outcome = self.outcome_list[i].copy()
                new_outcome_list.append(new_outcome)
        
        
        #Update attributes
        self.states = new_states
        self.cl_prob = new_cl_prob
        self.outcome_list = new_outcome_list
        
        
    #This functions performs measurements
    #The logic behind it is rather similar to reset, but this time we log the measurement outcome
    def measure(self, qubit_index, clbit_index):
        
        if thread_limit > 1:
            
            threads = []
            
            for i in range(len(self.states)):
                mp_wrapper = lambda : self.states[i].measure(qubit_index)
                thr = threading.Thread(target = mp_wrapper)
                
                if len(threads) > thread_limit:
                    threads[0].join()
                    threads.pop(0)
                
                
                thr.start()
                threads.append(thr)
            
        else:
            
            for i in range(len(self.states)):
                self.states[i].measure(qubit_index, True)

        
        new_states = []
        new_cl_prob = []
        new_outcome_list = []
        
        for i in range(len(self.states)):
            
            #Wait for measurement threads to finish
            if thread_limit > 1:
                if len(threads):
                    threads[0].join()
                    threads.pop(0)
            
            p_0, state_0, p_1, state_1 = self.states[i].last_mes_outcome
            
            if p_0 != 0:
                new_states.append(state_0)
                new_cl_prob.append(p_0*self.cl_prob[i])
                new_outcome = self.outcome_list[i].copy()
                
                #Here is the point, where we log the measurement outcome
                new_outcome[clbit_index] = 0
                new_outcome_list.append(new_outcome)
            if p_1 != 0:
                
                
                new_states.append(state_1)
                new_cl_prob.append(p_1*self.cl_prob[i])
                new_outcome = self.outcome_list[i].copy()
                new_outcome[clbit_index] = 1
                new_outcome_list.append(new_outcome)
                
        self.states = new_states
        self.cl_prob = new_cl_prob
        self.outcome_list = new_outcome_list
    
    
    def add_qubit(self):
        for state in self.states: state.add_qubit()
        
    def add_clbit(self):
        
        for i in range(len(self.outcome_list)):
            self.outcome_list[i] = np.array([0] + list(self.outcome_list[i]), dtype = np.byte)
