#!/usr/bin/env python3
"""Minimal educational reference: dimensionless, lossless, scalar two-beam model.

It is NOT a complete XRD simulator or a port of the MROX production kernel.
It does NOT include crystal-specific chi, geometry, absorption, instrumental
response or the MROXStudio 41 convention conversion.

Convention used:
    -i dU/dT = -V
    -i dV/dT = U + 2*eta*V
    X = V/U
    -i dX/dT = X**2 + 2*eta*X + 1.
"""
from __future__ import annotations
import numpy as np


def mobius_propagate(x_in: complex, eta: float, T: float) -> complex:
    """Exact 2x2 exponentiation (without scipy) for a constant eta and T."""
    s = np.sqrt(1.0 - complex(eta)**2)
    sh_over_s = T if abs(s) < 1e-10 else np.sinh(s*T)/s
    co = np.cosh(s*T)
    # exp(A*T) = exp(i*eta*T) * (co*I + sh_over_s*(A-i*eta*I)).
    # Global exp(i*eta*T) cancels in V/U.
    m11 = co - 1j*eta*sh_over_s
    m12 = -1j*sh_over_s
    m21 = 1j*sh_over_s
    m22 = co + 1j*eta*sh_over_s
    return (m21+m22*x_in)/(m11+m12*x_in)


def analytic_riccati(x_in: complex, eta: float, T: float) -> complex:
    """Closed-form Riccati propagator; teaching use at moderate |T| only."""
    xi = np.sqrt(complex(eta)**2-1)
    if abs(xi) < 1e-9:  # eta=+/-1: use matrix limit
        return mobius_propagate(x_in,eta,T)
    x_plus, x_minus = -eta+xi, -eta-xi
    K = (x_in-x_plus)/(x_in-x_minus)*np.exp(2j*xi*T)
    return (x_plus-K*x_minus)/(1-K)


def check():
    # Numerical equivalence in THIS convention; does not establish MROX equivalence.
    for eta in (-3.0,-1.1,-1.0,-.7,0,.7,1.0,1.1,3.0):
        for T in (.02,.2,.8,1.5,2.5):
            x1=mobius_propagate(0+0j,eta,T)
            x2=analytic_riccati(0+0j,eta,T)
            assert abs(x1-x2)<2e-10,(eta,T,x1,x2)
            assert abs(x1)**2<=1+1e-9
    # The zero-thickness limit must preserve an arbitrary amplitude ratio.
    for x0 in (0j,.2+.1j,-.6+.3j):
        assert abs(mobius_propagate(x0,.7,0)-x0)<1e-12
    print('PASS: analytic Riccati == 2x2 Möbius reference; T=0 identity')

if __name__ == '__main__': check()