//! # Quantum Mechanics
//!
//! This module provides symbolic tools for quantum mechanics, including representations
//! of quantum states (Bra-Ket notation), operators, and fundamental equations like
//! the time-independent and time-dependent Schrödinger equations. It also supports
//! concepts from perturbation theory and scattering processes.
use crate::symbolic::core::Expr;
use crate::symbolic::{calculus::differentiate, solve::solve};
use std::sync::Arc;
/// Represents a quantum state using Dirac notation (Ket).
#[derive(Clone, Debug)]
pub struct Ket {
    pub state: Expr,
}
/// Represents a quantum state using Dirac notation (Bra).
#[derive(Clone, Debug)]
pub struct Bra {
    pub state: Expr,
}
/// Computes the inner product of a Bra and a Ket, `<Bra|Ket>`.
///
/// This is a symbolic representation of the inner product over all space,
/// typically defined as `∫ ψ*(x)φ(x) dx`.
///
/// # Arguments
/// * `bra` - The `Bra` state `ψ*`.
/// * `ket` - The `Ket` state `φ`.
///
/// # Returns
/// An `Expr` representing the symbolic inner product.
pub fn bra_ket(bra: &Bra, ket: &Ket) -> Expr {
    Expr::Integral {
        integrand: Arc::new(Expr::new_mul(bra.state.clone(), ket.state.clone())),
        var: Arc::new(Expr::Variable("space".to_string())),
        lower_bound: Arc::new(Expr::NegativeInfinity),
        upper_bound: Arc::new(Expr::Infinity),
    }
}
/// Represents a quantum operator.
#[derive(Clone, Debug)]
pub struct Operator {
    pub op: Expr,
}
impl Operator {
    /// Applies an operator to a Ket, `O|Ket>`.
    ///
    /// # Arguments
    /// * `ket` - The `Ket` state to apply the operator to.
    ///
    /// # Returns
    /// A new `Ket` representing the transformed state.
    pub fn apply(&self, ket: &Ket) -> Ket {
        Ket {
            state: Expr::new_mul(self.op.clone(), ket.state.clone()),
        }
    }
}
/// Solves the time-independent Schrödinger equation `H|psi> = E|psi>`.
///
/// This function symbolically represents the solution of the eigenvalue problem
/// for the Hamiltonian operator `H` and the wave function `|psi>`.
///
/// # Arguments
/// * `hamiltonian` - The `Operator` representing the Hamiltonian `H`.
/// * `wave_function` - The `Ket` representing the wave function `|psi>`.
///
/// # Returns
/// A tuple `(eigenvalues, eigenfunctions)` where `eigenvalues` is a `Vec<Expr>`
/// and `eigenfunctions` is a `Vec<Ket>`.
pub fn solve_time_independent_schrodinger(
    hamiltonian: &Operator,
    wave_function: &Ket,
) -> (Vec<Expr>, Vec<Ket>) {
    let h_psi = hamiltonian.apply(wave_function);
    let e = Expr::Variable("E".to_string());
    let e_psi = Expr::new_mul(e, wave_function.state.clone());
    let equation = Expr::new_sub(h_psi.state, e_psi);
    let solutions = solve(&equation, "E");
    let eigenfunctions = solutions.iter().map(|_sol| wave_function.clone()).collect();
    (solutions, eigenfunctions)
}
/// Represents the time-dependent Schrödinger equation `i*hbar*d/dt|psi> = H|psi>`.
///
/// This equation describes how the quantum state of a physical system changes over time.
///
/// # Arguments
/// * `hamiltonian` - The `Operator` representing the Hamiltonian `H`.
/// * `wave_function` - The `Ket` representing the wave function `|psi>`.
///
/// # Returns
/// An `Expr` representing the symbolic time-dependent Schrödinger equation.
pub fn time_dependent_schrodinger_equation(hamiltonian: &Operator, wave_function: &Ket) -> Expr {
    let i = Expr::new_complex(Expr::Constant(0.0), Expr::Constant(1.0));
    let hbar = Expr::Variable("hbar".to_string());
    let i_hbar = Expr::new_mul(i, hbar);
    let d_psi_dt = differentiate(&wave_function.state, "t");
    let lhs = Expr::new_mul(i_hbar, d_psi_dt);
    let rhs = hamiltonian.apply(wave_function).state;
    Expr::new_sub(lhs, rhs)
}
/// Computes the first-order energy correction in perturbation theory.
///
/// In quantum mechanics, perturbation theory is a set of approximation schemes
/// related to a small disturbance applied to a system. The first-order energy
/// correction `E^(1)` is given by the expectation value of the perturbation `H'`
/// in the unperturbed state `|ψ^(0)>`: `E^(1) = <ψ^(0)|H'|ψ^(0)>`.
///
/// # Arguments
/// * `perturbation` - The `Operator` representing the perturbation `H'`.
/// * `unperturbed_state` - The `Ket` representing the unperturbed state `|ψ^(0)>`.
///
/// # Returns
/// An `Expr` representing the first-order energy correction.
pub fn first_order_energy_correction(perturbation: &Operator, unperturbed_state: &Ket) -> Expr {
    bra_ket(
        &Bra {
            state: unperturbed_state.state.clone(),
        },
        &perturbation.apply(unperturbed_state),
    )
}
/// Represents a scattering process.
///
/// This function symbolically represents the scattering amplitude, often using
/// approximations like the Born approximation. The scattering amplitude relates
/// the initial and final states of particles in a scattering event.
///
/// # Arguments
/// * `initial_state` - The `Ket` representing the initial state of the system.
/// * `final_state` - The `Ket` representing the final state of the system.
/// * `potential` - The `Operator` representing the scattering potential.
///
/// # Returns
/// An `Expr` representing the symbolic scattering amplitude.
pub fn scattering_amplitude(initial_state: &Ket, final_state: &Ket, potential: &Operator) -> Expr {
    let term = potential.apply(initial_state);
    bra_ket(
        &Bra {
            state: final_state.state.clone(),
        },
        &term,
    )
}
