#pragma once

#include "QuantumCircuitElement.h"

#include "QuantumWire.h"
#include <vector>

class QuantumGate : public QuantumCircuitElement
{
public:
    enum class GateType
    {
        X,
        Y,
        Z,
        H,
        CNOT,
        CUSTOM
    };
    virtual ~QuantumGate() = default;

    virtual std::size_t getNumInputs() const = 0;

    virtual std::size_t getNumOutputs() const = 0;

    virtual AbstractQuantumWire* getInput(std::size_t idx) const = 0;

    virtual AbstractQuantumWire* getOutput(std::size_t idx) const = 0;

    virtual GateType getGateType() const = 0;

    Type getType() const override
    {
        return Type::GATE;
    }
};
using QuantumGatePtr = std::unique_ptr<QuantumGate>;


class NInMOutQuantumGate : public QuantumGate
{
public:
    NInMOutQuantumGate(std::size_t numIn, std::size_t numOut, std::vector<AbstractQuantumWire*> inputs = {}, std::vector<AbstractQuantumWire*> outputs = {});

    virtual ~NInMOutQuantumGate() = default;

    std::size_t getNumInputs() const override;

    std::size_t getNumOutputs() const override;

    AbstractQuantumWire* getInput(std::size_t idx) const override;

    AbstractQuantumWire* getOutput(std::size_t idx) const override;

    void setAtInput(std::size_t idx, AbstractQuantumWire* value);

    void setAtOutput(std::size_t idx, AbstractQuantumWire* value);

private:
    std::size_t mNumIn{ 1 };
    std::size_t mNumOut{ 1 };

    std::vector<AbstractQuantumWire*> mInputs;
    std::vector<AbstractQuantumWire*> mOutputs;
};

class TwoInOneOutQuantumGate : public NInMOutQuantumGate
{
public:
    TwoInOneOutQuantumGate(AbstractQuantumWire* input0 = nullptr, AbstractQuantumWire* input1 = nullptr, AbstractQuantumWire* output = nullptr);

    void setInput0(AbstractQuantumWire* input);

    void setInput1(AbstractQuantumWire* input);

    void setOutput(AbstractQuantumWire* output);
};

class OneInOneOutQuantumGate : public NInMOutQuantumGate
{
public:
    OneInOneOutQuantumGate(AbstractQuantumWire* input = nullptr, AbstractQuantumWire* output = nullptr);

    void setInput(AbstractQuantumWire* input);

    void setOutput(AbstractQuantumWire* output);
};