using Random, LinearAlgebra

function generate_states(N, qubitsCount)
    dim = 2^qubitsCount  # Dimension of each state vector
    states = [normalize!(rand(dim)) for _ in 1:N]  # Generate and normalize N vectors
    return states
end

function read_operations(filename)
    open(filename, "r") do file
        lines = readlines(file)
        N, qubitsCount, S = parse.(Int, split(lines[1]))  # Read first line as numbers
        operations = [(parse(Int, split(line)[2])) for line in lines[2:end]]  # Extract qubit indices
        return N, qubitsCount, S, operations
    end
end

function pauli_x_matrix(qubitsCount, target_qubit)
    I = Matrix(LinearAlgebra.I, 2, 2)  # 2x2 Identity matrix
    X = [0 1; 1 0]  # Pauli-X matrix

    # Start with an identity matrix of correct size
    result = 1  # Start with scalar (not identity matrix)

    for i in 1:qubitsCount
        if i == target_qubit
            result = kron(result, X)  # Apply X gate
        else
            result = kron(result, I)  # Apply identity for other qubits
        end
    end
    return result
end


function apply_operations(states, qubitsCount, operations)
    for target_qubit in operations
        X_matrix = pauli_x_matrix(qubitsCount, target_qubit)
        states = [X_matrix * state for state in states]  # Apply transformation
    end
    return states
end

function measure_states(states)
    counts = Dict{String, Int}()  # Store counts of |00⟩, |01⟩, etc.

    for state in states
        max_index = argmax(abs.(state))  # Find highest probability component
        bitstring = string(Int(max_index) - 1, base=2)  # Convert to binary string
        counts[bitstring] = get(counts, bitstring, 0) + 1
    end
    return counts
end

using Base.Threads

function apply_operations_parallel(states, qubitsCount, operations)
    @threads for i in eachindex(states)
        for target_qubit in operations
            X_matrix = pauli_x_matrix(qubitsCount, target_qubit)
            states[i] = X_matrix * states[i]
        end
    end
    return states
end

using CUDA

function apply_operations_gpu(states, qubitsCount, operations)
    states_gpu = CuArray(hcat(states...))  # Move states to GPU
    for target_qubit in operations
        X_matrix = CuArray(pauli_x_matrix(qubitsCount, target_qubit))
		# Perform matrix multiplication on GPU
        states_gpu .= X_matrix * states_gpu
    end
    return Array(states_gpu)  # Move back to CPU
end

using BenchmarkTools

function benchmark_operations(states, qubitsCount, operations)
	# Benchmark the sequential implementation
	println("Benchmarking Sequential Implementation...")
	b = @benchmark apply_operations($states, $qubitsCount, $operations)
	display(b)
	# Benchmark the parallel implementation
	println("Benchmarking Parallel Implementation...")
	b = @benchmark apply_operations_parallel($states, $qubitsCount, $operations)
	display(b)
	# Benchmark the GPU implementation
	println("Benchmarking GPU Implementation...")
	b = @benchmark apply_operations_gpu($states, $qubitsCount, $operations)
	display(b)
end

function main(filename)
    N, qubitsCount, S, operations = read_operations(filename)
    states = generate_states(N, qubitsCount)

    # Sequential execution
    sequential_result = apply_operations(states, qubitsCount, operations)

    # Parallel execution
    parallel_result = apply_operations_parallel(states, qubitsCount, operations)

    # GPU execution
    gpu_result = apply_operations_gpu(states, qubitsCount, operations)

    println("Measurement Results (Sequential): ", measure_states(sequential_result))
    println("Measurement Results (Parallel): ", measure_states(parallel_result))
    println("Measurement Results (GPU): ", measure_states(gpu_result))

	# Benchmarking
	benchmark_operations(states, qubitsCount, operations)
end

# Run with an input file
main("input.txt")