"""
Lindblad master equation evolution for open quantum systems.

Implements the GKSL (Gorini-Kossakowski-Sudarshan-Lindblad) master equation:
    drho/dt = -i[H, rho] + sum_k gamma_k (L_k rho L_k^dag - 1/2 {L_k^dag L_k, rho})

Extracted and corrected from the v4 scientific implementation.
Key fix: fidelity is computed against a caller-supplied target state,
not against the maximally mixed state.

References:
    [1] Lindblad, Commun. Math. Phys. 48:119 (1976)
    [2] Gorini, Kossakowski, Sudarshan, J. Math. Phys. 17:821 (1976)
    [3] Breuer & Petruccione, "The Theory of Open Quantum Systems" (2002)
"""

import numpy as np
from scipy import linalg
from scipy.integrate import solve_ivp
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass, field

from .fidelity import (
    uhlmann_fidelity, von_neumann_entropy, purity,
    linear_entropy, validate_density_matrix
)
from .config import PHYS


@dataclass
class QuantumMeasures:
    """Computed quantum information measures at a timestep."""
    fidelity: float          # F(rho, target) -- against specified target
    von_neumann_entropy: float
    purity: float
    linear_entropy: float
    effective_dimension: float


@dataclass
class EvolutionResult:
    """Result of a Lindblad evolution run."""
    density_matrices: List[np.ndarray]
    measures: List[QuantumMeasures]
    times: np.ndarray


def build_lindblad_operators(dim: int) -> Tuple[List[np.ndarray], List[str]]:
    """
    Construct standard Lindblad operators for a d-dimensional system.

    Returns:
        operators: List of Lindblad operator matrices
        labels: Human-readable labels for each operator
    """
    operators = []
    labels = []

    # Amplitude damping: |i><i+1| (energy lowering)
    for i in range(dim - 1):
        L = np.zeros((dim, dim), dtype=complex)
        L[i, i + 1] = 1.0
        operators.append(L)
        labels.append(f'amp_damp_{i}_{i+1}')

    # Phase damping: |i><i| (diagonal dephasing)
    for i in range(dim):
        L = np.zeros((dim, dim), dtype=complex)
        L[i, i] = 1.0
        operators.append(L)
        labels.append(f'phase_damp_{i}')

    # Depolarizing: generalized Pauli-like operators
    if dim >= 2:
        # X-type (cyclic permutation)
        L_x = np.zeros((dim, dim), dtype=complex)
        for i in range(dim):
            L_x[(i + 1) % dim, i] = 1.0
        operators.append(L_x)
        labels.append('depol_X')

        # Z-type (phase)
        L_z = np.diag([np.exp(2j * np.pi * i / dim) for i in range(dim)])
        operators.append(L_z)
        labels.append('depol_Z')

    return operators, labels


def compute_decoherence_rates(
    dim: int,
    system_params: Dict[str, float],
) -> List[float]:
    """
    Compute Lindblad decoherence rates from experimental T1, T2, gate fidelity.

    Physical rates:
        - Amplitude damping:  gamma_ad = 1/T1
        - Phase damping:      gamma_pd = 1/T2 - 1/(2*T1)  (pure dephasing)
        - Depolarizing:       gamma_dep = (1 - F_gate) / (d^2 - 1)

    Args:
        dim: Hilbert space dimension
        system_params: Dict with 'T1', 'T2', 'gate_fidelity' keys
    """
    T1 = system_params.get('T1', PHYS.ibm_q_params['T1'])
    T2 = system_params.get('T2', PHYS.ibm_q_params['T2'])
    gate_fid = system_params.get('gate_fidelity', 0.999)

    gamma_ad = 1.0 / T1
    gamma_pd = max(1.0 / T2 - 1.0 / (2.0 * T1), 0.0)
    gamma_dep = (1.0 - gate_fid) / max(dim * dim - 1, 1)

    rates = []
    rates.extend([gamma_ad] * (dim - 1))  # amplitude damping operators
    rates.extend([gamma_pd] * dim)         # phase damping operators
    if dim >= 2:
        rates.extend([gamma_dep] * 2)      # depolarizing (X, Z)

    return rates


def normalize_density_matrix(rho: np.ndarray) -> np.ndarray:
    """
    Project a matrix onto the space of valid density matrices.

    Enforces:
        1. Hermiticity: rho = rho^dagger
        2. Positive semidefinite: all eigenvalues >= 0
        3. Trace = 1
    """
    # Hermiticity
    rho = (rho + rho.conj().T) / 2.0

    # Positive semidefinite via eigendecomposition
    eigvals, eigvecs = linalg.eigh(rho)
    eigvals = np.maximum(eigvals, 0.0)

    # Reconstruct and normalize trace
    rho = eigvecs @ np.diag(eigvals) @ eigvecs.conj().T
    trace = np.real(np.trace(rho))
    if trace > 1e-15:
        rho = rho / trace

    return rho


class LindbladEvolver:
    """
    Evolves a density matrix under the Lindblad master equation.

    drho/dt = -i[H, rho] + sum_k gamma_k (L_k rho L_k^dag - 1/2 {L_k^dag L_k, rho})
    """

    def __init__(
        self,
        hamiltonian: np.ndarray,
        lindblad_ops: List[np.ndarray],
        rates: List[float],
    ):
        """
        Args:
            hamiltonian: System Hamiltonian (d x d, Hermitian)
            lindblad_ops: List of Lindblad operators
            rates: Decoherence rate for each operator
        """
        assert len(lindblad_ops) == len(rates), \
            f"Number of operators ({len(lindblad_ops)}) must match rates ({len(rates)})"

        self.dim = hamiltonian.shape[0]
        self.H = hamiltonian
        self.L_ops = lindblad_ops
        self.rates = rates

        # Precompute L_dag_L for each operator
        self._L_dag_L = [L.conj().T @ L for L in lindblad_ops]

    def lindblad_rhs(self, t: float, rho_vec: np.ndarray) -> np.ndarray:
        """
        RHS of the Lindblad equation in vectorized form for ODE integration.

        Args:
            t: Time (unused for time-independent H, but required by solve_ivp)
            rho_vec: Flattened density matrix

        Returns:
            Flattened time derivative drho/dt
        """
        rho = rho_vec.reshape((self.dim, self.dim))
        drho = np.zeros_like(rho)

        # Unitary part: -i[H, rho]
        drho += -1j * (self.H @ rho - rho @ self.H)

        # Dissipator: sum_k gamma_k (L_k rho L_k^dag - 1/2 {L_k^dag L_k, rho})
        for L, gamma, LdL in zip(self.L_ops, self.rates, self._L_dag_L):
            if gamma <= 0:
                continue
            L_dag = L.conj().T
            drho += gamma * (L @ rho @ L_dag - 0.5 * (LdL @ rho + rho @ LdL))

        return drho.flatten()

    def evolve(
        self,
        rho_init: np.ndarray,
        dt: float,
        n_steps: int,
        target_state: np.ndarray,
        method: str = 'RK45',
    ) -> EvolutionResult:
        """
        Evolve density matrix and compute measures at each step.

        Args:
            rho_init: Initial density matrix
            dt: Time step
            n_steps: Number of evolution steps
            target_state: Reference state for fidelity computation.
                          Typically the ideal unitary-evolved state.
            method: ODE integration method

        Returns:
            EvolutionResult with full history
        """
        rho = rho_init.copy()
        density_matrices = [rho.copy()]
        measures = [self._compute_measures(rho, target_state)]
        times = [0.0]

        for step in range(n_steps):
            rho_vec = rho.flatten()

            sol = solve_ivp(
                self.lindblad_rhs,
                [0, dt],
                rho_vec,
                method=method,
                rtol=1e-8,
                atol=1e-10,
            )

            rho = sol.y[:, -1].reshape((self.dim, self.dim))
            rho = normalize_density_matrix(rho)

            t = (step + 1) * dt
            density_matrices.append(rho.copy())
            measures.append(self._compute_measures(rho, target_state))
            times.append(t)

        return EvolutionResult(
            density_matrices=density_matrices,
            measures=measures,
            times=np.array(times),
        )

    def evolve_single_step(
        self,
        rho: np.ndarray,
        dt: float,
        method: str = 'RK45',
    ) -> np.ndarray:
        """Evolve density matrix by one time step. Returns new rho."""
        rho_vec = rho.flatten()
        sol = solve_ivp(
            self.lindblad_rhs,
            [0, dt],
            rho_vec,
            method=method,
            rtol=1e-8,
            atol=1e-10,
        )
        rho_new = sol.y[:, -1].reshape((self.dim, self.dim))
        return normalize_density_matrix(rho_new)

    def _compute_measures(
        self,
        rho: np.ndarray,
        target_state: np.ndarray,
    ) -> QuantumMeasures:
        """Compute quantum information measures against the target state."""
        return QuantumMeasures(
            fidelity=uhlmann_fidelity(rho, target_state),
            von_neumann_entropy=von_neumann_entropy(rho),
            purity=purity(rho),
            linear_entropy=linear_entropy(rho),
            effective_dimension=1.0 / max(purity(rho), 1e-15),
        )


def unitary_evolve(rho: np.ndarray, H: np.ndarray, t: float) -> np.ndarray:
    """
    Ideal unitary evolution: rho(t) = U rho(0) U^dag where U = exp(-iHt).

    Use this to compute the target state for fidelity comparison:
    the state the system SHOULD be in if there were no decoherence.
    """
    U = linalg.expm(-1j * H * t)
    return U @ rho @ U.conj().T