#include <iostream>
#include <vector>
#include <unordered_map>
#include <iomanip>
#include <algorithm>
#include <unordered_set>
#include <queue>
#include <random>

#define CNOT_VALUE 1.0
#define SWAP_VALUE 3.0*CNOT_VALUE
#define ERROR_SCALE 0.0001

struct VirtualQubit {
    int id;
    double importance = 0.0; // accumulates from gates
};

struct Gate {
    std::string type; // "X", "CNOT", etc
    std::vector<int> qubits; // 1 or 2 qubit ids
    double weight = 1.0; // weight of this gate
};

class VirtualCircuit {
	public:
		std::unordered_map<int, VirtualQubit> vqubits;
		std::vector<Gate> gates;
		

		void add_qubit(int id) {
			vqubits.emplace(id, VirtualQubit{id});
		}

		void add_gate(const std::string& type, const std::vector<int>& qubits, double weight=1.0) {
			gates.push_back(Gate{type, qubits, weight});

			// Update importance factor for virtual qubits
			for(auto qid : qubits) {
				vqubits.at(qid).importance += weight;
			}
		}

		void print_circuit() const {
			std::cout << "\nVirtual Circuit:\n";
			for(auto& [id, vq] : vqubits)
				std::cout << "VQubit " << id << " | Importance: " << vq.importance << "\n";

			std::cout << "\nGates:\n";
			for(auto& g : gates) {
				std::cout << g.type << " on qubits: ";
				for(auto qid : g.qubits) std::cout << qid << " ";
				std::cout << " | Weight: " << g.weight << "\n";
			}
		}
};


class Qubit {
    public:
        int id;
        double error_rate;
        
        
        VirtualQubit* mapped_vqubit = nullptr;
        
        double getError() const{
            return error_rate;
        }
        
        //Err 0.1 importance: 3.5
        //1/Err * importance
        double getCost() const{
            if(mapped_vqubit == nullptr){
                return 0.0;
            }
            return mapped_vqubit->importance * (ERROR_SCALE / error_rate);
        }
    
        // IDs of qubits this qubit can directly interact with
        std::vector<int> connections;
    
        Qubit(int id, double error_rate): id(id), error_rate(error_rate) {}
    
        void add_connection(int other_id) {
            connections.push_back(other_id);
        }
};

class QuantumDevice {
    private:
        std::unordered_map<int, Qubit> qubits;
        std::unordered_map<int, int> vqubit_to_pqubit;
    
    public:
        void add_qubit(int id, double error_rate) {
            qubits.emplace(id, Qubit(id, error_rate));
        }
        
        int optimization = 0;
        void apply_circuit(const VirtualCircuit& vc,bool verbose = true) {
            std::vector<Qubit*> physicalOrder;
            for (auto& [id, q] : qubits) {
                physicalOrder.push_back(&q);
            }
            
            std::vector<const VirtualQubit*> virtualOrder;
            for (auto& [id, vq] : vc.vqubits) {
                virtualOrder.push_back(&vq);
            }
        
            if (optimization > 0) { //Opt1
                //collect all qubits
                std::vector<Qubit*> allQubits;
                for (auto& [id, q] : qubits) {
                    allQubits.push_back(&q);
                }
                
                //sort by error (ascending)
                std::sort(allQubits.begin(), allQubits.end(), [](const Qubit* a, const Qubit* b) {
                    return a->error_rate < b->error_rate;
                });
                
                //separate worst 5 qubits
                int worstCount = std::min(5, (int)allQubits.size());
                std::unordered_set<int> worstIds;
                
                for (int i = (int)allQubits.size() - worstCount; i < (int)allQubits.size(); i++) {
                    worstIds.insert(allQubits[i]->id);
                }
                
                //priority BFS (min-heap based on error_rate)
                std::vector<Qubit*> ordered;
                std::unordered_set<int> visited;
                
                // min-heap: smallest error_rate first
                auto cmp = [&](int a, int b) {
                    return qubits.at(a).error_rate > qubits.at(b).error_rate;
                };
				
                std::priority_queue<int, std::vector<int>, decltype(cmp)> pq(cmp);
                
                // start from best qubit
                pq.push(allQubits[0]->id);
                visited.insert(allQubits[0]->id);
                
                while (!pq.empty()) {
                    int current = pq.top();
                    pq.pop();
                
                    if (worstIds.count(current)) continue;
                
                    ordered.push_back(&qubits.at(current));
                
                    for (int neighbor : qubits.at(current).connections) {
                        if (!visited.count(neighbor)) {
                            visited.insert(neighbor);
                            pq.push(neighbor);
                        }
                    }
                }
                
                //add any remaining non-worst (disconnected components)
                for (auto* qb : allQubits) {
                    if (!visited.count(qb->id) && !worstIds.count(qb->id)) {
                        ordered.push_back(qb);
                    }
                }
                
                //append worst qubits at the end (still sorted among themselves)
                for (auto* qb : allQubits) {
                    if (worstIds.count(qb->id)) {
                        ordered.push_back(qb);
                    }
                }
                
                // Final assignment
                physicalOrder = ordered;
                
                //Virtual circuit
                std::unordered_map<int, std::vector<int>> vConnections;

                for (const auto& gate : vc.gates) {
                    if (gate.qubits.size() == 2) {
                        int a = gate.qubits[0];
                        int b = gate.qubits[1];
                        vConnections[a].push_back(b);
                        vConnections[b].push_back(a);
                    }
                }
                
                std::vector<const VirtualQubit*> orderedVirtual;
                
                std::unordered_set<int> visitedV;
                
                // Start from most important qubit
                int startId = virtualOrder[0]->id;
                
                // max-heap (higher importance first)
                auto vcmp = [&](int a, int b) {
                    return vc.vqubits.at(a).importance < vc.vqubits.at(b).importance;
                };
                
                std::priority_queue<int, std::vector<int>, decltype(vcmp)> vpq(vcmp);
                
                vpq.push(startId);
                visitedV.insert(startId);
                
                while (!vpq.empty()) {
                    int current = vpq.top();
                    vpq.pop();
                
                    orderedVirtual.push_back(&vc.vqubits.at(current));
                
                    for (int neighbor : vConnections[current]) {
                        if (!visitedV.count(neighbor)) {
                            visitedV.insert(neighbor);
                            vpq.push(neighbor);
                        }
                    }
                }
                
                // Handle disconnected virtual qubits
                for (auto& [id, vq] : vc.vqubits) {
                    if (!visitedV.count(id)) {
                        orderedVirtual.push_back(&vq);
                    }
                }
                
                virtualOrder = orderedVirtual;
               
            }
        
            // Store mapping: virtual qubit id -> physical qubit id
            vqubit_to_pqubit.clear();
            size_t mappingCount = std::min(virtualOrder.size(), physicalOrder.size());
            for (size_t i = 0; i < mappingCount; i++) {
                Qubit* pq = physicalOrder[i];
                pq->mapped_vqubit = const_cast<VirtualQubit*>(virtualOrder[i]);
                vqubit_to_pqubit[virtualOrder[i]->id] = pq->id;
            }
        
            printEvaledCost(vc,verbose);
        }
        
        void printEvaledCost(const VirtualCircuit& vc,bool verbose) {
            if(verbose)
                std::cout << "Cost of each qubit operations (if cost is high, good):" << std::endl;
        
            double totalCost = 0.0;
            for (auto it = qubits.begin(); it != qubits.end(); ++it) {
                double tcost = it->second.getCost();
                if (tcost != 0.0) {
                    totalCost += tcost;
                    if(verbose)
                        std::cout << "Qubit number " << it->second.id << ": " << tcost << std::endl;
                }
            }
        
            // For each 2-qubit gate, check if physical qubits are connected via DFS
            // If not connected, penalize by 2*SWAP_VALUE / edge_error
            if(verbose) std::cout << "\nSWAP Penalty Analysis:\n";
            double totalPenalty = 0.0;
        
            for (auto& gate : vc.gates) {
                if (gate.qubits.size() != 2) continue;
        
                int vid1 = gate.qubits[0];
                int vid2 = gate.qubits[1];
        
                // Look up physical qubit ids for these virtual qubits
                if (vqubit_to_pqubit.find(vid1) == vqubit_to_pqubit.end()) continue;
                if (vqubit_to_pqubit.find(vid2) == vqubit_to_pqubit.end()) continue;
        
                int pid1 = vqubit_to_pqubit[vid1];
                int pid2 = vqubit_to_pqubit[vid2];
        
                // DFS to check connectivity between pid1 and pid2
                // and find the path's worst (highest) edge error
                std::unordered_map<int, bool> visited;
                std::unordered_map<int, int> parent;  // to trace path
                std::vector<int> stack = {pid1};
                visited[pid1] = true;
                bool connected = false;
        
                while (!stack.empty() && !connected) {
                    int cur = stack.back();
                    stack.pop_back();
        
                    if (cur == pid2) {
                        connected = true;
                        break;
                    }
        
                    for (int neighbor : qubits.at(cur).connections) {
                        if (!visited[neighbor]) {
                            visited[neighbor] = true;
                            parent[neighbor] = cur;
                            stack.push_back(neighbor);
                        }
                    }
                }
        
                if (!connected) {
                    if(verbose) std::cout << "  VQubit " << vid1 << "->" << vid2
                              << " (PQubit " << pid1 << "->" << pid2 << "): NOT REACHABLE\n";
                    continue;
                }
        
                // Trace path from pid2 back to pid1
                std::vector<int> path;
                for (int cur = pid2; cur != pid1; cur = parent[cur]) {
                    path.push_back(cur);
                }
                path.push_back(pid1);
                std::reverse(path.begin(), path.end());
        
                // Count hops that are NOT direct neighbors (each non-direct hop = 1 SWAP needed)
                // Here every hop in the path beyond the first is a real edge, so SWAPs = path.size()-2
                // (direct connection = path size 2, no SWAPs needed)
                int swapsNeeded = (int)path.size() - 2;
        
                if (swapsNeeded <= 0) {
                    if(verbose) std::cout << "  VQubit " << vid1 << "->" << vid2
                              << " (PQubit " << pid1 << "->" << pid2 << "): directly connected, no SWAP needed\n";
                    continue;
                }
        
                // Use the average error rate along the path edges as edge_error
                double edgeErrorSum = 0.0;
                for (int node : path) {
                    edgeErrorSum += qubits.at(node).error_rate;
                }
                double edgeError = edgeErrorSum / path.size();
        
                double penalty = swapsNeeded * gate.weight *SWAP_VALUE * ((ERROR_SCALE/edgeError));
                totalPenalty += penalty;
        
                if(verbose) std::cout << "  VQubit " << vid1 << "->" << vid2
                          << " (PQubit " << pid1 << "->" << pid2 << ")"
                          << " | Hops: " << path.size() - 1
                          << " | SWAPs needed: " << swapsNeeded
                          << " | Edge error: " << edgeError
                          << " | Penalty: " << penalty << "\n";
            }
        
            if(verbose) std::cout << "\nTotal gain before penalty: " << totalCost << std::endl;
            if(verbose) std::cout << "Total SWAP penalty:        " << totalPenalty << std::endl;
            if(verbose) std::cout << "Net cost effectivity:      " << totalCost - totalPenalty << std::endl;
            n += 1.0;
            AtotalCost+=totalCost;
            AtotalPenalty+=totalPenalty;
        }
        
        double n = 0.0;
        double AtotalCost = 0.0;
        double AtotalPenalty = 0.0;
        
        void printAVGValues() const{
            if(n == 0.0){
                std::cout << "No valid placements happened for AVG data" << std::endl;
                return;
            }
            std::cout << "\nAVG. Total gain before penalty: " << AtotalCost/n << std::endl;
            std::cout << "AVG. Total SWAP penalty:        " << AtotalPenalty/n << std::endl;
            std::cout << "AVG. Net cost effectivity:      " << AtotalCost/n - AtotalPenalty/n << std::endl;
        }
    
        // Bidirectional connection (like simplified real hardware coupling maps)
        void connect_qubits(int id1, int id2) {
			qubits.at(id1).add_connection(id2);
			qubits.at(id2).add_connection(id1);
		}
		
    	Qubit& getQubitWithMostError() {
            auto it = qubits.begin();
        
            if (it == qubits.end()) {
                throw std::runtime_error("Empty circuit");
            }
        
            auto maxIt = it;
            ++it;
        
            for (; it != qubits.end(); ++it) {
                if (it->second.getError() > maxIt->second.getError()) {
                    maxIt = it;
                }
            }
        
            return maxIt->second;
        }
        
        Qubit& getQubitWithLeastError() {
            auto it = qubits.begin();
        
            if (it == qubits.end()) {
                throw std::runtime_error("Empty circuit");
            }
        
            auto minIt = it;
            ++it;
        
            for (; it != qubits.end(); ++it) {
                if (it->second.getError() < minIt->second.getError()) {
                    minIt = it;
                }
            }
        
            return minIt->second;
        }
        
    
        void print_device() const {
            for (const auto& [id, q] : qubits) {
                std::cout << "Qubit " << q.id
                          << " | Error rate: " << std::fixed << std::setprecision(4)
                          << q.error_rate << "\n";
    
                std::cout << "  Connected to: ";
                for (int conn : q.connections) {
                    std::cout << conn << " ";
                }
                std::cout << "\n\n";
            }
        }
    
        // Check if a 2-qubit gate is allowed
        bool can_apply_two_qubit_gate(int id1, int id2) const {
            const auto& q = qubits.at(id1);
            for (int conn : q.connections) {
                if (conn == id2) {
                    return true;
                }
            }
            return false;
        }
        
        void initFalcon(){
            for (int i = 0; i <= 26; i++) {
                add_qubit(i, ERROR_SCALE + ERROR_SCALE * ((i + 1 ) % 10));
            }
        
            connect_qubits(1, 0);
            connect_qubits(4, 1);
            connect_qubits(7, 4);
            connect_qubits(10, 7);
            connect_qubits(12, 10);
        
            connect_qubits(15, 18);
            connect_qubits(18, 21);
            connect_qubits(21, 23);
        
            connect_qubits(5, 3);
            connect_qubits(8, 5);
            connect_qubits(11, 8);
            connect_qubits(14, 11);
        
            connect_qubits(16, 19);
            connect_qubits(19, 22);
            connect_qubits(22, 25);
            connect_qubits(25, 26);
        
            connect_qubits(7, 6);
            connect_qubits(18, 17);
        
            connect_qubits(8, 9);
            connect_qubits(19, 20);
            connect_qubits(23, 24);
        
            connect_qubits(13, 14);
            connect_qubits(13, 12);
        
            connect_qubits(2, 3);
            connect_qubits(1, 2);
        
            connect_qubits(24, 25);
        
            connect_qubits(12, 15);
            connect_qubits(14, 16);
        }
};

VirtualCircuit generate_random_circuit(
    int numQubits,
    int numGates,
    double twoQubitProb = 0.7,
    double longRangeProb = 0.3
) {
    VirtualCircuit vc;

    // Create qubits
    for (int i = 0; i < numQubits; i++) {
        vc.add_qubit(i);
    }

    std::random_device rd;
    std::mt19937 gen(rd());

    std::uniform_real_distribution<> probDist(0.0, 1.0);
    std::uniform_int_distribution<> qubitDist(0, numQubits - 1);
    std::uniform_real_distribution<> weightDist(0.5, 3.0);

    for (int i = 0; i < numGates; i++) {
        double weight = weightDist(gen);

        // Decide gate type
        if (probDist(gen) < twoQubitProb) {
            int q1 = qubitDist(gen);
            int q2;

            if (probDist(gen) < longRangeProb) {
                // Force long-range interaction
                do {
                    q2 = qubitDist(gen);
                } while (q2 == q1);
            } else {
                // Prefer nearby indices (pseudo-locality)
                int offset = (gen() % 3) - 1; // -1, 0, +1
                q2 = std::clamp(q1 + offset, 0, numQubits - 1);
                if (q2 == q1) q2 = (q1 + 1) % numQubits;
            }

            vc.add_gate("CNOT", {q1, q2}, weight);
        } else {
            int q = qubitDist(gen);

            if (probDist(gen) < 0.5)
                vc.add_gate("H", {q}, weight * 0.5);
            else
                vc.add_gate("X", {q}, weight * 0.5);
        }
    }

    return vc;
}

int main() {
    QuantumDevice device;
    
    device.initFalcon();
    
    device.optimization = 0; //0 for naive, 1 for heuristic
    const uint8_t SAMPLE_ID = 1; //Which test to run.
    
    
    VirtualCircuit vc;
    
    switch(SAMPLE_ID){
        case 1: //Sparse graph and local chain structures
            for(int i = 0; i < 10; i++) {
                vc.add_qubit(i);
            }
            
            //Local chains
            vc.add_gate("CNOT", {0,1}, 1.0);
            vc.add_gate("CNOT", {1,2}, 1.0);
            vc.add_gate("CNOT", {2,3}, 1.0);
            vc.add_gate("CNOT", {3,4}, 1.0);
            
            vc.add_gate("CNOT", {5,6}, 1.0);
            vc.add_gate("CNOT", {6,7}, 1.0);
            vc.add_gate("CNOT", {7,8}, 1.0);
            vc.add_gate("CNOT", {8,9}, 1.0);
            
            //Cross-cluster entanglement (forces routing)
            vc.add_gate("CNOT", {0,5}, 2.0);
            vc.add_gate("CNOT", {2,7}, 2.0);
            vc.add_gate("CNOT", {4,9}, 2.0);
            
            //Star pattern (hotspot qubit)
            vc.add_gate("CNOT", {1,6}, 1.5);
            vc.add_gate("CNOT", {1,7}, 1.5);
            vc.add_gate("CNOT", {1,8}, 1.5);
            
            //Long-range conflicts
            vc.add_gate("CNOT", {0,9}, 3.0);
            vc.add_gate("CNOT", {3,8}, 2.5);
            
            //Single qubit noise (importance skew)
            vc.add_gate("H", {0}, 0.5);
            vc.add_gate("X", {1}, 0.5);
            vc.add_gate("H", {2}, 0.5);
            vc.add_gate("X", {7}, 0.5);
            vc.add_gate("H", {9}, 0.5);
            
            //Dense subgraph (hard case)
            vc.add_gate("CNOT", {2,5}, 2.0);
            vc.add_gate("CNOT", {2,6}, 2.0);
            vc.add_gate("CNOT", {3,5}, 2.0);
            vc.add_gate("CNOT", {4,6}, 2.0);
            
            //Reuse + interference
            vc.add_gate("CNOT", {6,0}, 2.5);
            vc.add_gate("CNOT", {7,3}, 2.5);
            vc.add_gate("CNOT", {8,1}, 2.5);
            
            device.apply_circuit(vc);
            break;
        case 2: //Random
            for(size_t i=0;i<1000;i++){
                
            
                vc = generate_random_circuit(
                    15,     // number of qubits
                    50,     // number of gates
                    0.5,   // % of 2-qubit gates
                    0.4     // % of long-range interactions
                );
                
                device.apply_circuit(vc,false);
                
                
            
            }
            device.printAVGValues();
            
            break;
        case 3: //Dense graph
            for(int i=0; i<7; i++) {
                vc.add_qubit(i);
            }
            for(int i = 0; i < 6; i++) {
                for(int j = i+1; j < 6; j++) {
                    vc.add_gate("CNOT", {i, j}, 1.5);
                }
            }
            
            device.apply_circuit(vc);
            break;
        case 4: //Simple sample
            for(int i=0; i<5; i++) {
                vc.add_qubit(i);
            }
            
            vc.add_gate("X", {0}, 1.0);
            vc.add_gate("CNOT", {0,1}, CNOT_VALUE);
            vc.add_gate("CNOT", {1,2}, CNOT_VALUE);
            vc.add_gate("H", {3}, 0.5);
            vc.add_gate("CNOT", {3,4}, CNOT_VALUE);
            
            device.apply_circuit(vc);
            
            break;
    }

    //vc.print_circuit();
    
    

    //device.print_device();

    return 0;
}