"""
vqe_solver.py

This module contains a function to solve a quantum operator using the
Variational Quantum Eigensolver (VQE). The VQE algorithm combines quantum and
classical optimization to find the minimum eigenvalue of a Hamiltonian.
"""

import numpy as np
from qiskit.circuit.library import EfficientSU2
from qiskit.quantum_info import Statevector, state_fidelity
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit_aer.backends.aer_simulator import AerBackend
from qiskit_ibm_runtime import EstimatorV2, Session
from scipy.optimize import minimize

from quantum_pipeline.circuits import HFData, build_hf_initial_state
from quantum_pipeline.configs.module.backend import BackendConfig
from quantum_pipeline.mappers.mapper import Mapper
from quantum_pipeline.solvers.optimizer_config import get_optimizer_configuration
from quantum_pipeline.solvers.solver import Solver
from quantum_pipeline.structures.vqe_observation import (
    VQEInitialData,
    VQEProcess,
    VQEResult,
)
from quantum_pipeline.utils.timer import Timer


class MaxFunctionEvalsReachedError(Exception):
    """Raised when the hard function evaluation limit is exceeded."""

    def __init__(self, iteration: int, best_params: np.ndarray, best_energy: float):
        super().__init__(
            f'Hard function evaluation limit reached after {iteration} evaluations.'
        )
        self.iteration = iteration
        self.best_params = best_params
        self.best_energy = best_energy


class VQESolver(Solver):
    def __init__(
        self,
        qubit_op,
        backend_config: BackendConfig,
        max_iterations=50,
        optimizer='COBYLA',
        ansatz_reps=3,
        default_shots=1024,
        convergence_threshold=None,
        optimization_level=3,
        seed=None,
        init_strategy='random',
        hf_data: HFData | None = None,
        mapper: Mapper | None = None,
    ):
        super().__init__()
        self.qubit_op = qubit_op
        self.ansatz_reps = ansatz_reps
        self.optimizer = optimizer
        self.max_iterations = max_iterations
        self.seed = seed
        self.digits_iter = len(str(max_iterations))
        self.default_shots = default_shots
        self.backend_config = backend_config
        self.vqe_process: list[VQEProcess] = []
        self.current_iter = 1
        self.convergence_threshold = convergence_threshold
        self.optimization_level = optimization_level
        self.init_strategy = init_strategy
        self.hf_data = hf_data
        self.mapper = mapper

    def _optimize_circuits(self, ansatz, hamiltonian, backend):
        """Prepare ISA-compatible circuits and observables"""
        target = backend.target
        pm = generate_preset_pass_manager(
            target=target, optimization_level=self.optimization_level
        )
        ansatz_isa = pm.run(ansatz)

        hamiltonian_isa = hamiltonian.apply_layout(layout=ansatz_isa.layout)
        return ansatz_isa, hamiltonian_isa

    def _build_ansatz(self, n_qubits):
        """Build the EfficientSU2 ansatz.

        Always builds a plain EfficientSU2 (no HF circuit prepended).
        When using HF init, we compute initial parameters that prepare the HF
        state through the ansatz instead — avoiding the problem where fixed CX
        entangling gates destroy the HF state at zero parameters.
        """
        return EfficientSU2(n_qubits, reps=self.ansatz_reps)

    def _compute_hf_initial_parameters(self, ansatz):
        """Find EfficientSU2 parameters that prepare the HF state.

        Performs a short classical pre-optimization to find parameters where
        the ansatz output matches the HF state (maximizes state fidelity).
        This avoids the problem of prepending HF circuit + zero params, where
        the fixed CX gates in EfficientSU2 destroy the HF state.
        """
        hf_circuit = build_hf_initial_state(self.hf_data, self.mapper)
        target_sv = Statevector(hf_circuit)

        def neg_fidelity(params):
            bound = ansatz.assign_parameters(params)
            sv = Statevector(bound)
            return -state_fidelity(sv, target_sv)

        best_fid = 0.0
        best_params = np.zeros(ansatz.num_parameters)
        n_attempts = 10
        seed = self.seed if self.seed is not None else 0

        for i in range(n_attempts):
            np.random.seed(seed + i)
            x0 = 2 * np.pi * np.random.random(ansatz.num_parameters)
            res = minimize(neg_fidelity, x0, method='COBYLA', options={'maxiter': 1000})
            fid = -res.fun
            if fid > best_fid:
                best_fid = fid
                best_params = res.x
            if best_fid > 0.9999:
                break

        self.logger.info(
            f'HF parameter pre-optimization: fidelity={best_fid:.6f} '
            f'after {i + 1} attempts'
        )
        return best_params

    def _compute_initial_parameters(self, ansatz):
        """Compute initial ansatz parameters based on the configured strategy."""
        param_num = ansatz.num_parameters

        if self.init_strategy == 'hf' and self.hf_data is not None and self.mapper is not None:
            self.logger.info('Computing HF-equivalent initial parameters...')
            return self._compute_hf_initial_parameters(ansatz)

        if self.init_strategy == 'hf' and (self.hf_data is None or self.mapper is None):
            self.logger.warning(
                'HF init strategy requested but no HF data/mapper available, falling back to random'
            )

        if self.seed is not None:
            np.random.seed(self.seed)
            self.logger.info(f'Using seed {self.seed} for parameter initialization')
        return 2 * np.pi * np.random.random(param_num)

    @property
    def _nuclear_repulsion(self) -> np.float64 | None:
        if self.hf_data is not None and self.hf_data.nuclear_repulsion_energy is not None:
            return np.float64(self.hf_data.nuclear_repulsion_energy)
        return None

    def _log_total_energy(self, electronic_energy: float) -> None:
        """Log total energy (electronic + nuclear repulsion) if nuclear repulsion is available."""
        if self._nuclear_repulsion is not None:
            total = electronic_energy + self._nuclear_repulsion
            self.logger.info(
                f'Total energy (electronic + nuclear repulsion): {total:.8f} Ha '
                f'(nuclear repulsion: {self._nuclear_repulsion:.8f} Ha)'
            )

    def _make_truncated_result(self, best_energy, best_params, elapsed):
        """Build a VQEResult from the best observed state after early termination."""
        return VQEResult(
            initial_data=self.init_data,
            iteration_list=self.vqe_process,
            minimum=np.float64(best_energy),
            optimal_parameters=best_params,
            maxcv=None,
            minimization_time=np.float64(elapsed),
            nuclear_repulsion_energy=self._nuclear_repulsion,
        )

    def compute_energy(self, params, ansatz, hamiltonian, estimator):
        """Return estimate of energy from estimator"""
        if self.max_iterations is not None and self.current_iter > self.max_iterations:
            best = min(self.vqe_process, key=lambda p: p.cumulative_min_energy)
            raise MaxFunctionEvalsReachedError(
                iteration=self.current_iter - 1,
                best_params=best.parameters,
                best_energy=float(best.cumulative_min_energy),
            )

        pub = (ansatz, [hamiltonian], [params])
        result = estimator.run(pubs=[pub]).result()
        energy, std = result[0].data.evs[0], result[0].data.stds[0]

        if self.vqe_process:
            prev = self.vqe_process[-1]
            energy_delta = np.float64(energy - prev.result)
            parameter_delta_norm = np.float64(np.linalg.norm(params - prev.parameters))
            cumulative_min_energy = np.float64(min(energy, prev.cumulative_min_energy))
        else:
            energy_delta = None
            parameter_delta_norm = None
            cumulative_min_energy = np.float64(energy)

        iter = VQEProcess(
            iteration=self.current_iter,
            parameters=params,
            result=energy,
            std=std,
            energy_delta=energy_delta,
            parameter_delta_norm=parameter_delta_norm,
            cumulative_min_energy=cumulative_min_energy,
        )

        self.vqe_process.append(iter)
        self.logger.debug(
            f'Iters. done: {self.current_iter:0{self.digits_iter}d} [Current cost: {energy}]'
        )
        self.current_iter += 1
        return energy

    def via_ibmq(self, backend):
        """Run the VQE simulation on IBM Quantum backend."""
        hamiltonian = self.qubit_op

        self.logger.info(f'Initializing the ansatz with {self.ansatz_reps} reps...')
        ansatz = self._build_ansatz(hamiltonian.num_qubits)
        self.logger.info('Ansatz initialized.')

        x0 = self._compute_initial_parameters(ansatz)
        self.logger.debug(f'Initial ansatz parameters:\n\n{x0}\n')

        self.logger.info('Optimizing ansatz and hamiltonian...')
        ansatz_isa, hamiltonian_isa = self._optimize_circuits(ansatz, hamiltonian, backend)
        self.logger.info('Ansatz and hamiltonian optimized.')

        self.init_data = VQEInitialData(
            backend=backend.name,
            num_qubits=hamiltonian_isa.num_qubits,
            hamiltonian=hamiltonian_isa.to_list(),
            num_parameters=ansatz_isa.num_parameters,
            initial_parameters=x0,
            optimizer=self.optimizer,
            ansatz=ansatz_isa,
            ansatz_reps=self.ansatz_reps,
            noise_backend=self.backend_config.noise if self.backend_config.noise else 'undef',
            default_shots=self.default_shots,
            seed=self.seed,
            init_strategy=self.init_strategy,
        )
        self.logger.info('Opening a session...')

        with Session(backend=backend) as session:
            estimator = EstimatorV2(mode=session)
            estimator.options.default_shots = self.default_shots

            optimization_params, minimize_tol = get_optimizer_configuration(
                optimizer=self.optimizer,
                max_iterations=self.max_iterations,
                convergence_threshold=self.convergence_threshold,
                num_parameters=len(x0),
            )

            self.logger.debug(f'Optimization params: {optimization_params}')
            if minimize_tol is not None:
                self.logger.debug(f'Minimize tolerance: {minimize_tol}')

            if self.convergence_threshold and self.max_iterations:
                self.logger.info(
                    f'Starting VQE optimization with max iterations {self.max_iterations} taking priority over '
                    f'convergence threshold {self.convergence_threshold}'
                )
            elif self.convergence_threshold:
                self.logger.info(
                    f'Starting VQE optimization with convergence threshold {self.convergence_threshold}'
                )
            elif self.max_iterations:
                self.logger.info(
                    f'Starting VQE optimization with max iterations {self.max_iterations}'
                )
            else:
                self.logger.info('Starting VQE optimization with default settings')

        truncated = None
        res = None
        with Timer() as t:
            try:
                res = minimize(
                    self.compute_energy,
                    x0,
                    args=(ansatz_isa, hamiltonian_isa, estimator),
                    method=self.optimizer,
                    options=optimization_params,
                    tol=minimize_tol,
                )
            except MaxFunctionEvalsReachedError as e:
                truncated = e

        if truncated is not None:
            self.logger.info(
                f'Hard evaluation limit reached after {truncated.iteration} function evaluations '
                f'(limit: {self.max_iterations}). Best energy: {truncated.best_energy:.8f} Ha'
            )
            self._log_total_energy(truncated.best_energy)
            return self._make_truncated_result(truncated.best_energy, truncated.best_params, t.elapsed)

        actual_iterations = len(self.vqe_process)
        best_energy = (
            min(p.cumulative_min_energy for p in self.vqe_process)
            if self.vqe_process
            else res.fun
        )
        if self.convergence_threshold:
            if res.success:
                self.logger.info(
                    f'VQE converged after {actual_iterations} iterations '
                    f'(threshold: {self.convergence_threshold}). '
                    f'Best energy: {best_energy:.8f} Ha'
                )
            else:
                self.logger.info(
                    f'VQE stopped after {actual_iterations} iterations - convergence not achieved. '
                    f'Best energy: {best_energy:.8f} Ha'
                )
        else:
            self.logger.info(
                f'VQE completed {actual_iterations} iterations. Best energy: {best_energy:.8f} Ha'
            )
        self._log_total_energy(best_energy)

        result = VQEResult(
            initial_data=self.init_data,
            iteration_list=self.vqe_process,
            minimum=res.fun,
            optimal_parameters=res.x,
            maxcv=getattr(res, 'maxcv', None),
            minimization_time=np.float64(t.elapsed),
            nuclear_repulsion_energy=self._nuclear_repulsion,
        )

        self.logger.info('Session closed.')
        self.logger.info(
            f'Calculations on quantum hardware via IBMQ completed in {t.elapsed:.6f} seconds.'
        )
        return result

    def via_aer(self, backend):
        """Run the VQE simulation via Aer simulator."""
        hamiltonian = self.qubit_op

        self.logger.info(f'Initializing the ansatz with {self.ansatz_reps} reps...')
        ansatz = self._build_ansatz(hamiltonian.num_qubits)
        self.logger.info('Ansatz initialized.')

        x0 = self._compute_initial_parameters(ansatz)
        self.logger.debug(f'Initial ansatz parameters:\n\n{x0}\n')

        self.logger.info('Optimizing ansatz and hamiltonian...')
        ansatz_isa, hamiltonian_isa = self._optimize_circuits(ansatz, hamiltonian, backend)
        self.logger.info('Ansatz and hamiltonian optimized.')

        self.init_data = VQEInitialData(
            backend=backend.name,
            num_qubits=hamiltonian_isa.num_qubits,
            hamiltonian=hamiltonian_isa.to_list(),
            num_parameters=ansatz_isa.num_parameters,
            initial_parameters=x0,
            optimizer=self.optimizer,
            ansatz=ansatz_isa,
            ansatz_reps=self.ansatz_reps,
            noise_backend=self.backend_config.noise if self.backend_config.noise else 'undef',
            default_shots=self.default_shots,
            seed=self.seed,
            init_strategy=self.init_strategy,
        )

        estimator = EstimatorV2(mode=backend)
        estimator.options.default_shots = self.default_shots

        optimization_params, minimize_tol = get_optimizer_configuration(
            optimizer=self.optimizer,
            max_iterations=self.max_iterations,
            convergence_threshold=self.convergence_threshold,
            num_parameters=len(x0),
        )

        self.logger.debug(f'Optimization params: {optimization_params}')
        if minimize_tol is not None:
            self.logger.debug(f'Minimize tolerance: {minimize_tol}')

        if self.convergence_threshold and self.max_iterations:
            self.logger.info(
                f'Starting VQE optimization with max iterations {self.max_iterations} taking priority over '
                f'convergence threshold {self.convergence_threshold}'
            )
        elif self.convergence_threshold:
            self.logger.info(
                f'Starting VQE optimization with convergence threshold {self.convergence_threshold}'
            )
        elif self.max_iterations:
            self.logger.info(
                f'Starting VQE optimization with max iterations {self.max_iterations}'
            )
        else:
            self.logger.info('Starting VQE optimization with default settings')
        truncated = None
        res = None
        with Timer() as t:
            try:
                res = minimize(
                    self.compute_energy,
                    x0,
                    args=(ansatz_isa, hamiltonian_isa, estimator),
                    method=self.optimizer,
                    options=optimization_params,
                    tol=minimize_tol,
                )
            except MaxFunctionEvalsReachedError as e:
                truncated = e

        if truncated is not None:
            self.logger.info(
                f'Hard evaluation limit reached after {truncated.iteration} function evaluations '
                f'(limit: {self.max_iterations}). Best energy: {truncated.best_energy:.8f} Ha'
            )
            self._log_total_energy(truncated.best_energy)
            return self._make_truncated_result(truncated.best_energy, truncated.best_params, t.elapsed)

        actual_iterations = len(self.vqe_process)
        best_energy = (
            min(p.cumulative_min_energy for p in self.vqe_process)
            if self.vqe_process
            else res.fun
        )
        if self.convergence_threshold:
            if res.success:
                self.logger.info(
                    f'VQE converged after {actual_iterations} iterations '
                    f'(threshold: {self.convergence_threshold}). '
                    f'Best energy: {best_energy:.8f} Ha'
                )
            else:
                self.logger.info(
                    f'VQE stopped after {actual_iterations} iterations - convergence not achieved. '
                    f'Best energy: {best_energy:.8f} Ha'
                )
        else:
            self.logger.info(
                f'VQE completed {actual_iterations} iterations. Best energy: {best_energy:.8f} Ha'
            )
        self._log_total_energy(best_energy)

        result = VQEResult(
            initial_data=self.init_data,
            iteration_list=self.vqe_process,
            minimum=res.fun,
            optimal_parameters=res.x,
            maxcv=getattr(res, 'maxcv', None),
            minimization_time=np.float64(t.elapsed),
            nuclear_repulsion_energy=self._nuclear_repulsion,
        )

        self.logger.info(
            f'Simulation via Aer completed in {t.elapsed:.6f} seconds and {len(result.iteration_list)} iterations.'
        )
        return result

    def solve(self):
        """Run the VQE simulation and return the result."""
        self.current_iter = 1
        self.vqe_process = []

        backend = self.get_backend()

        return self.via_aer(backend) if isinstance(backend, AerBackend) else self.via_ibmq(backend)
