
import numpy as np
from .grover             import grover_search
from .qaoa               import run_qaoa
from .vqe                import run_vqe, run_adapt_vqe, run_sampling_vqe
from .qpe                import qpe, iterative_qpe, hamiltonian_qpe
from .amplitude_estimation import (amplitude_estimation,
                                    iterative_amplitude_estimation,
                                    mle_amplitude_estimation)
from .time_evolution     import trotter_evolution, varqrte, varqite
from .qml                import build_qnn, train_qnn, quantum_kernel, quantum_kernel_matrix
from .qft                import qft, inverse_qft
from .state_prep         import basis_encoding, amplitude_encoding, angle_encoding
from .variational        import (parameter_shift_gradient,
                                  optimize_parameters, variational_loop)

class AdvancedQuantum:
    """
    PKTron v2.3.2+ — Unified Advanced Algorithms API
    One class, all algorithms. Designed for Google Colab + research workflows.
    """

    # ── Grover ──────────────────────────────────────────────────────────────
    def grover(self, n_qubits: int = 3, target: int = 5) -> dict:
        """Grover search. Returns measured state, probabilities, success flag."""
        return grover_search(n_qubits, target)

    # ── QAOA ────────────────────────────────────────────────────────────────
    def qaoa(self, problem: dict = None, p_layers: int = 1) -> dict:
        """QAOA for MaxCut. problem = {'n_nodes': int, 'edges': list of tuples}"""
        if problem is None:
            problem = {"n_nodes": 4, "edges": [(0,1),(1,2),(2,3),(3,0)]}
        return run_qaoa(problem, p_layers)

    # ── VQE ─────────────────────────────────────────────────────────────────
    def vqe(self, hamiltonian: np.ndarray = None, n_qubits: int = 2,
            mode: str = "standard") -> dict:
        """VQE variants: mode = standard | adapt | sampling"""
        if hamiltonian is None:
            hamiltonian = np.diag([-1.0, 0.0, 0.0, 1.0])  # simple 2-qubit diagonal H
        if mode == "adapt":
            return run_adapt_vqe(hamiltonian, n_qubits)
        elif mode == "sampling":
            return run_sampling_vqe(hamiltonian, n_qubits)
        else:
            return run_vqe(hamiltonian, n_qubits)

    # ── QPE ─────────────────────────────────────────────────────────────────
    def qpe(self, unitary: np.ndarray = None, eigenstate: np.ndarray = None,
            n_ancilla: int = 4, mode: str = "standard") -> dict:
        """QPE variants: mode = standard | iterative | hamiltonian"""
        if unitary is None:
            phase   = 0.375
            unitary = np.array([[np.exp(2j*np.pi*phase), 0],
                                 [0, np.exp(-2j*np.pi*phase)]])
        if eigenstate is None:
            eigenstate = np.array([1.0, 0.0], dtype=complex)
        if mode == "iterative":
            return iterative_qpe(unitary, eigenstate, n_ancilla)
        elif mode == "hamiltonian":
            H = np.array([[-1, 0],[0, 1]], dtype=complex) * 0.5
            return hamiltonian_qpe(H, eigenstate, n_ancilla)
        else:
            return qpe(unitary, eigenstate, n_ancilla)

    # ── Amplitude Estimation ─────────────────────────────────────────────────
    def amplitude_estimation(self, a: float = 0.3, n_ancilla: int = 4,
                              mode: str = "standard") -> dict:
        """AE variants: mode = standard | iterative | mle"""
        fn = lambda: a
        if mode == "iterative":
            return iterative_amplitude_estimation(fn)
        elif mode == "mle":
            return mle_amplitude_estimation(fn)
        else:
            return amplitude_estimation(fn, n_ancilla)

    # ── Time Evolution ───────────────────────────────────────────────────────
    def time_evolution(self, hamiltonian: np.ndarray = None,
                       initial_state: np.ndarray = None,
                       time: float = 1.0, mode: str = "trotter") -> dict:
        """Time evolution: mode = trotter | varqrte | varqite"""
        if hamiltonian is None:
            hamiltonian = np.array([[0, 1],[1, 0]], dtype=complex) * 0.5
        if initial_state is None:
            initial_state = np.array([1.0, 0.0], dtype=complex)
        if mode == "varqrte":
            return varqrte(hamiltonian, initial_state, time)
        elif mode == "varqite":
            return varqite(hamiltonian, initial_state, time)
        else:
            return trotter_evolution(hamiltonian, initial_state, time)

    # ── QNN ─────────────────────────────────────────────────────────────────
    def qnn(self, data: np.ndarray = None, labels: np.ndarray = None,
            n_qubits: int = 2, mode: str = "train") -> dict:
        """QNN: mode = train | kernel"""
        if data is None:
            data   = np.random.uniform(0, np.pi, (6, n_qubits))
            labels = np.array([1,-1,1,-1,1,-1], dtype=float)
        if mode == "kernel":
            K = quantum_kernel_matrix(data, n_qubits)
            return {"kernel_matrix": K, "n_samples": len(data)}
        else:
            return train_qnn(data, labels, n_qubits=n_qubits)

    # ── QFT ─────────────────────────────────────────────────────────────────
    def qft(self, n_qubits: int = 3, state: np.ndarray = None,
            inverse: bool = False) -> dict:
        """Forward or inverse QFT."""
        if state is None:
            dim   = 2 ** n_qubits
            state = np.eye(dim)[0].astype(complex)  # |0...0>
        return inverse_qft(n_qubits, state) if inverse else qft(n_qubits, state)

    # ── State Prep ───────────────────────────────────────────────────────────
    def state_prep(self, data, mode: str = "amplitude") -> dict:
        """State preparation: mode = basis | amplitude | angle"""
        data = np.array(data)
        if mode == "basis":
            return basis_encoding(data.astype(int).tolist())
        elif mode == "angle":
            return angle_encoding(data)
        else:
            return amplitude_encoding(data)

    # ── Variational Tools ────────────────────────────────────────────────────
    def variational(self, cost_fn=None, n_params: int = 4,
                    mode: str = "optimize") -> dict:
        """Variational framework: mode = optimize | loop | gradient"""
        if cost_fn is None:
            target = np.random.uniform(0, np.pi, n_params)
            cost_fn = lambda p: float(np.sum((p - target)**2))
        init = np.random.uniform(-np.pi, np.pi, n_params)
        if mode == "gradient":
            g = parameter_shift_gradient(cost_fn, init)
            return {"gradient": g, "params": init}
        elif mode == "loop":
            return variational_loop(cost_fn, init)
        else:
            return optimize_parameters(cost_fn, init)
