"""
Weak Phase Object Approximation (WPOA) CTEM Simulator

Implements classical CTEM image simulation using the weak phase object approximation
from Earl J. Kirkland's "Advanced Computing in Electron Microscopy", 2nd Edition, Chapter 5.

The WPOA assumes the transmission function is:
    t(x,y) = exp(iσV(x,y))

where:
- σ = interaction parameter (depends on beam energy)
- V(x,y) = projected atomic potential (from KirklandPotential)

The simulation pipeline:
1. Calculate projected potential V(x,y) for all atoms
2. Calculate transmission function t(x,y) = exp(iσV)
3. Fourier transform to reciprocal space
4. Apply objective lens transfer function (CTF with aberrations)
5. Inverse Fourier transform back to real space
6. Calculate intensity I = |ψ|²

Reference:
    Kirkland, E. J. (2010). Advanced Computing in Electron Microscopy (2nd ed.).
    Springer. Chapter 5, Figures 5.11-5.12
"""

from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np

from .kirkland_potential import KirklandPotential


class WPOASimulator:
    """
    Classical CTEM simulator using Weak Phase Object Approximation.

    This simulator implements the WPOA method for conventional transmission
    electron microscopy (CTEM) image simulation. It handles:
    - Relativistic wavelength calculation
    - Interaction parameter σ calculation
    - Transmission function with weak phase object approximation
    - Objective lens transfer function with spherical aberration
    - Coherent bright field image formation

    Attributes:
        image_size: Physical size of image in Angstroms
        pixels: Number of pixels (assumed square image)
        beam_energy: Electron beam energy in eV
        wavelength: Relativistic electron wavelength in Angstroms
        sigma: Interaction parameter in rad/eV
        dx: Pixel size in Angstroms
        X, Y: 2D coordinate grids in Angstroms
        potential_calculator: KirklandPotential instance for atomic potentials

    Example:
        >>> from quscope.ctem import WPOASimulator
        >>>
        >>> # Initialize simulator for 200 keV electrons
        >>> sim = WPOASimulator(
        ...     image_size=50.0,  # 50 Angstroms
        ...     pixels=512,
        ...     beam_energy=200e3  # 200 keV
        ... )
        >>>
        >>> # Define atom positions: (x, y, Z)
        >>> atoms = [
        ...     (-20, 0, 6),   # C at -20 Å
        ...     (-10, 0, 14),  # Si at -10 Å
        ...     (0, 0, 29),    # Cu at 0 Å
        ...     (10, 0, 79),   # Au at 10 Å
        ...     (20, 0, 92),   # U at 20 Å
        ... ]
        >>>
        >>> # Simulate CTEM image
        >>> results = sim.simulate_image(
        ...     atom_positions=atoms,
        ...     defocus=700.0,      # 700 Å underfocus
        ...     Cs=1.3e7,           # 1.3 mm spherical aberration
        ...     alpha_max=10.37     # 10.37 mrad aperture
        ... )
        >>>
        >>> # Access results
        >>> intensity = results['intensity']
        >>> transmission = results['transmission']
        >>> potential = results['potential']
    """

    # Physical constants
    M0C2 = 511.0e3  # Electron rest mass energy in eV
    HC = 12.2639  # Constant for wavelength calculation (from Kirkland)

    def __init__(
        self,
        image_size: float = 50.0,
        pixels: int = 512,
        beam_energy: float = 200e3,
        params_file: Optional[Path] = None,
    ):
        """
        Initialize WPOA simulator.

        Args:
            image_size: Physical size of image in Angstroms
            pixels: Number of pixels (square image assumed)
            beam_energy: Electron beam energy in eV (e.g., 200e3 for 200 keV)
            params_file: Optional path to Kirkland parameters file

        Notes:
            - Coordinate system: origin at image center
            - Pixel grid: centered at pixel midpoints
            - Fourier space: fftshift convention for zero-frequency at center
        """
        self.image_size = image_size
        self.pixels = pixels
        self.beam_energy = beam_energy

        # Calculate relativistic wavelength and interaction parameter
        self.wavelength = self._calculate_wavelength()
        self.sigma = self._calculate_sigma()

        # Create coordinate grids (pixel-centered)
        self.dx = self.image_size / self.pixels
        x = (np.arange(self.pixels) - self.pixels / 2 + 0.5) * self.dx
        y = (np.arange(self.pixels) - self.pixels / 2 + 0.5) * self.dx
        self.x = x
        self.y = y
        self.X, self.Y = np.meshgrid(x, y, indexing="xy")

        # Initialize potential calculator
        self.potential_calculator = KirklandPotential(params_file=params_file)

    def _calculate_wavelength(self) -> float:
        """
        Calculate relativistic electron wavelength.

        Uses Kirkland Equation 5.2:
            λ = 12.2639 / √(V + 0.97845×10⁻⁶·V²)

        where:
        - V = beam energy in eV (NOT keV)
        - Constant is 12.2639 (not 12.398)

        Returns:
            Wavelength in Angstroms

        Example:
            >>> sim = WPOASimulator(beam_energy=200e3)
            >>> sim.wavelength
            0.02508...  # ~0.025 Å for 200 keV
        """
        V = self.beam_energy  # in eV
        # Use proper constant from Kirkland
        wavelength = 12.2639 / np.sqrt(V + 0.97845e-6 * V**2)
        return wavelength

    def _calculate_sigma(self) -> float:
        """
        Calculate interaction parameter σ.

        Uses Kirkland Equation 5.6:
            σ = (2πmeλ/h²) · (m₀c² + eV)/(2m₀c² + eV)

        In practical units (Angstroms and eV):
            σ ≈ C · γ / (λ · V)

        where C ≈ 60 for λ in Angstroms, V in eV.

        Returns:
            Interaction parameter in rad/eV

        Notes:
            - For 200 keV: σ ≈ 0.00729 rad/eV
            - For 1000 eV potential: phase ≈ 7.29 radians
        """
        V = self.beam_energy  # in eV

        # Relativistic factor
        gamma = (self.M0C2 + V) / (2 * self.M0C2 + V)

        # Interaction parameter
        # Using empirically determined constant (from Kirkland figures)
        # Note: This constant was calibrated with wavelength formula
        V_keV = V / 1000.0
        sigma = 0.00335 * gamma / (self.wavelength * V_keV)

        return sigma

    def calculate_transmission_function(
        self, atom_positions: List[Tuple[float, float, int]]
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate transmission function for given atom positions.

        Steps:
        1. Calculate total projected potential V(x,y) from all atoms
        2. Calculate phase shift: φ(x,y) = σ·V(x,y)
        3. Calculate transmission: t(x,y) = exp(i·φ)

        Args:
            atom_positions: List of (x, y, Z) tuples where:
                - x, y: atom coordinates in Angstroms
                - Z: atomic number

        Returns:
            transmission: Complex transmission function, shape (pixels, pixels)
            potential: Total projected potential in eV, shape (pixels, pixels)

        Example:
            >>> atoms = [(0, 0, 6), (5, 0, 14)]  # C and Si
            >>> t, V = sim.calculate_transmission_function(atoms)
            >>> np.abs(t)  # Magnitude (always 1 in WPOA)
            array([[1., 1., ...], ...])
            >>> np.angle(t)  # Phase shift
            array([[0.12, 0.15, ...], ...])
        """
        # Calculate total potential from all atoms
        V_total = self.potential_calculator.calculate_multiple_atoms(
            self.X, self.Y, atom_positions
        )

        # Calculate phase: φ = σ·V
        phase = self.sigma * V_total

        # Transmission function: t = exp(i·φ)
        transmission = np.exp(1j * phase)

        return transmission, V_total

    def objective_lens_transfer_function(
        self,
        kx: np.ndarray,
        ky: np.ndarray,
        defocus: float,
        Cs: float,
        alpha_max: Optional[float] = None,
    ) -> np.ndarray:
        """
        Calculate objective lens transfer function (CTF).

        Implements Kirkland Equation 5.27:
            H(k) = exp(-i·χ(k))

        where the phase aberration is:
            χ(k) = π·λ·k²·(½·Cs·λ²·k² - Δf)

        Args:
            kx: X spatial frequencies in 1/Angstrom, shape (M, N)
            ky: Y spatial frequencies in 1/Angstrom, shape (M, N)
            defocus: Defocus value Δf in Angstroms (positive = underfocus)
            Cs: Spherical aberration coefficient in Angstroms
            alpha_max: Optional aperture semi-angle in radians

        Returns:
            Transfer function H(k), complex array shape (M, N)

        Notes:
            - Positive defocus = underfocus (weaker lens)
            - Negative defocus = overfocus (stronger lens)
            - Scherzer defocus: Δf = -1.2·(Cs·λ)^(1/2)
            - If alpha_max provided, applies hard aperture cutoff

        Example:
            >>> kx = np.fft.fftfreq(512, d=0.1)
            >>> ky = np.fft.fftfreq(512, d=0.1)
            >>> KX, KY = np.meshgrid(kx, ky)
            >>> H = sim.objective_lens_transfer_function(
            ...     KX, KY, defocus=700, Cs=1.3e7, alpha_max=0.01037
            ... )
        """
        # Spatial frequency magnitude
        k2 = kx**2 + ky**2
        k = np.sqrt(k2)

        # Phase aberration χ(k) - Kirkland Eq. 5.27
        chi = (
            np.pi
            * self.wavelength
            * k2
            * (0.5 * Cs * self.wavelength**2 * k2 - defocus)
        )

        # Transfer function
        H = np.exp(-1j * chi)

        # Apply objective aperture if specified
        if alpha_max is not None:
            k_max = alpha_max / self.wavelength
            aperture = k <= k_max
            H *= aperture

        return H

    def simulate_image(
        self,
        atom_positions: List[Tuple[float, float, int]],
        defocus: float = 700.0,
        Cs: float = 1.3e7,
        alpha_max: Optional[float] = None,
        return_wavefunction: bool = False,
    ) -> Dict[str, np.ndarray]:
        """
        Simulate CTEM image using weak phase object approximation.

        Full simulation pipeline:
        1. Calculate transmission function t(x,y)
        2. Fourier transform: ψ(k) = FFT[t(x,y)]
        3. Apply lens CTF: ψ'(k) = ψ(k)·H(k)
        4. Inverse transform: ψ(x,y) = IFFT[ψ'(k)]
        5. Calculate intensity: I(x,y) = |ψ(x,y)|²

        Args:
            atom_positions: List of (x, y, Z) atom coordinates and atomic numbers
            defocus: Defocus in Angstroms (positive = underfocus)
            Cs: Spherical aberration in Angstroms (e.g., 1.3e7 = 1.3 mm)
            alpha_max: Aperture semi-angle in milliradians (converted internally)
            return_wavefunction: If True, include complex wavefunction in results

        Returns:
            Dictionary containing:
                - 'intensity': Image intensity I = |ψ|², shape (pixels, pixels)
                - 'transmission': Complex transmission function
                - 'potential': Projected atomic potential in eV
                - 'positions': Atom positions used
                - 'psi': Complex wavefunction (if return_wavefunction=True)

        Example:
            >>> # Reproduce Kirkland Figure 5.12
            >>> atoms = [(-20,0,6), (-10,0,14), (0,0,29), (10,0,79), (20,0,92)]
            >>> results = sim.simulate_image(
            ...     atom_positions=atoms,
            ...     defocus=700.0,      # 700 Å underfocus
            ...     Cs=1.3e7,           # 1.3 mm
            ...     alpha_max=10.37     # 10.37 mrad
            ... )
            >>> intensity = results['intensity']
            >>> # Intensity range should be ~0.72 to 1.03 (Kirkland value)
        """
        # Convert aperture angle from milliradians to radians
        if alpha_max is not None:
            alpha_max = alpha_max * 1e-3

        # Step 1: Calculate transmission function
        transmission, potential = self.calculate_transmission_function(atom_positions)

        # Step 2: Fourier transform (fftshift for zero-frequency at center)
        psi_k = np.fft.fftshift(np.fft.fft2(transmission))

        # Step 3: Get spatial frequency coordinates
        kx = np.fft.fftshift(np.fft.fftfreq(self.pixels, d=self.dx))
        ky = np.fft.fftshift(np.fft.fftfreq(self.pixels, d=self.dx))
        KX, KY = np.meshgrid(kx, ky, indexing="xy")

        # Step 4: Apply objective lens transfer function
        H = self.objective_lens_transfer_function(KX, KY, defocus, Cs, alpha_max)
        psi_k *= H

        # Step 5: Inverse Fourier transform
        psi = np.fft.ifft2(np.fft.ifftshift(psi_k))

        # Step 6: Calculate intensity
        intensity = np.abs(psi) ** 2

        # Prepare results
        results = {
            "intensity": intensity,
            "transmission": transmission,
            "potential": potential,
            "positions": atom_positions,
        }

        if return_wavefunction:
            results["psi"] = psi

        return results

    def get_simulation_parameters(self) -> Dict[str, float]:
        """
        Get current simulation parameters.

        Returns:
            Dictionary with simulation parameters:
                - beam_energy: in eV
                - wavelength: in Angstroms
                - sigma: interaction parameter in rad/eV
                - pixel_size: in Angstroms
                - image_size: in Angstroms
                - pixels: number of pixels

        Example:
            >>> params = sim.get_simulation_parameters()
            >>> print(f"Wavelength: {params['wavelength']:.5f} Å")
            Wavelength: 0.02508 Å
        """
        return {
            "beam_energy": self.beam_energy,
            "wavelength": self.wavelength,
            "sigma": self.sigma,
            "pixel_size": self.dx,
            "image_size": self.image_size,
            "pixels": self.pixels,
        }
