from __future__ import annotations

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
radial_schrodinger_solver.py

実用寄りの 1D 動径 Schrödinger 方程式ソルバ
- 任意の有効ポテンシャル Veff(r) を関数で指定可能
- 左右からの shooting + 中点接続（対数微分一致）
- Euler / Verlet / Numerov を比較可能
- 水素様原子では解析解との比較も可能

方程式（原子単位系）:
    u''(r) = 2 [Veff(r) - E] u(r)

ここで
    u(r) = r R(r)

水素様原子では
    Veff(r) = -Z/r + l(l+1)/(2r^2)

なので
    u'' = [l(l+1)/r^2 - 2Z/r - 2E] u
になる。

使い方例:
    python H_like_wf.py --potential hydrogenic --Z 1 --n 1 --l 0 --solve-e --method all --plot
    python H_like_wf.py --potential hydrogenic --Z 1 --n 1 --l 0 --solve-e --method numerov --plot
    python H_like_wf.py --potential hydrogenic --Z 1 --n 2 --l 0 --solve-e --method all --plot
    python H_like_wf.py --potential hydrogenic --Z 1 --n 3 --l 1 --solve-e --method numerov --plot

"""

import argparse
import math
from dataclasses import dataclass
from typing import Callable, Optional, Tuple, Dict, List

import numpy as np
import matplotlib.pyplot as plt


# ============================================================
# 基本データ構造
# ============================================================

@dataclass
class PotentialParams:
    potential_name: str = "hydrogenic"
    Z: float = 1.0
    l: int = 0
    soft_core_a: float = 0.0   # 将来用: -Z/sqrt(r^2+a^2) などに拡張可能


@dataclass
class SolveParams:
    rmin: float = 1.0e-5
    rmax: float = 50.0
    nmesh: int = 6000
    match_fraction: float = 0.35   # 左から何割地点を接続点にするか
    method: str = "numerov"
    solve_e: bool = True
    E_guess: Optional[float] = None
    scan_rel_width: float = 0.4
    scan_points: int = 500
    bisect_tol: float = 1.0e-12
    bisect_maxiter: int = 120


@dataclass
class SolveResult:
    method: str
    E: float
    r: np.ndarray
    u: np.ndarray
    R: np.ndarray
    nodes: int
    match_index: int
    mismatch: float


# ============================================================
# ポテンシャル
# ============================================================

def Veff_hydrogenic(r: np.ndarray | float, p: PotentialParams) -> np.ndarray | float:
    """
    水素様原子の有効ポテンシャル
        Veff(r) = -Z/r + l(l+1)/(2r^2)
    """
    Z = p.Z
    l = p.l
    return -Z / r + l * (l + 1) / (2.0 * r * r)


def get_potential_function(name: str) -> Callable[[np.ndarray | float, PotentialParams], np.ndarray | float]:
    if name == "hydrogenic":
        return Veff_hydrogenic
    raise ValueError(f"Unknown potential: {name}")


def g_from_Veff(r: float, E: float, p: PotentialParams, Veff_func: Callable) -> float:
    """
    u'' = g(r) u
    g(r) = 2 [Veff(r) - E]
    """
    return 2.0 * (Veff_func(r, p) - E)


# ============================================================
# 水素様原子の原点近傍級数
# ============================================================

def hydrogenic_origin_series_coeffs(Z: float, E: float, l: int) -> Tuple[float, float]:
    """
    水素様原子での
        u(r)=r^(l+1)(1+a1 r+a2 r^2+...)
    の係数
    """
    a1 = -Z / (l + 1.0)
    a2 = (2.0 * Z * Z / (l + 1.0) - 2.0 * E) / (4.0 * l + 6.0)
    return a1, a2


def hydrogenic_initial_u_up(r0: float, Z: float, E: float, l: int) -> Tuple[float, float]:
    a1, a2 = hydrogenic_origin_series_coeffs(Z, E, l)
    p = l + 1
    s = 1.0 + a1 * r0 + a2 * r0 * r0
    sp = a1 + 2.0 * a2 * r0
    u0 = (r0 ** p) * s
    up0 = p * (r0 ** (p - 1)) * s + (r0 ** p) * sp
    return u0, up0


def right_tail_initial_u_up(r1: float, E: float, p: PotentialParams, Veff_func: Callable) -> Tuple[float, float]:
    """
    右端での減衰初期条件。
    束縛状態で Veff(rmax) > E を仮定し
        u ~ exp(-kappa r)
    と置く簡易版。
    """
    kappa2 = 2.0 * (Veff_func(r1, p) - E)
    if kappa2 <= 0:
        # 右端でまだ禁制領域でない場合の保険
        kappa = 1.0e-6
    else:
        kappa = math.sqrt(kappa2)

    # 非常に小さい値で規格化
    u1 = math.exp(-kappa * r1)
    up1 = -kappa * u1
    return u1, up1


# ============================================================
# 各積分法: 片側積分
# ============================================================

def integrate_left_euler(r: np.ndarray, E: float, p: PotentialParams, Veff_func: Callable) -> np.ndarray:
    h = r[1] - r[0]
    u = np.zeros_like(r)
    v = np.zeros_like(r)

    if p.potential_name == "hydrogenic":
        u[0], v[0] = hydrogenic_initial_u_up(r[0], p.Z, E, p.l)
    else:
        raise NotImplementedError("一般ポテンシャルの原点条件は個別実装してください。")

    for i in range(len(r) - 1):
        g = g_from_Veff(r[i], E, p, Veff_func)
        u[i + 1] = u[i] + h * v[i]
        v[i + 1] = v[i] + h * g * u[i]
    return u


def integrate_right_euler(r: np.ndarray, E: float, p: PotentialParams, Veff_func: Callable) -> np.ndarray:
    h = r[1] - r[0]
    u = np.zeros_like(r)
    v = np.zeros_like(r)

    u[-1], v[-1] = right_tail_initial_u_up(r[-1], E, p, Veff_func)

    for i in range(len(r) - 1, 0, -1):
        g = g_from_Veff(r[i], E, p, Veff_func)
        u[i - 1] = u[i] - h * v[i]
        v[i - 1] = v[i] - h * g * u[i]
    return u


def integrate_left_verlet(r: np.ndarray, E: float, p: PotentialParams, Veff_func: Callable) -> np.ndarray:
    h = r[1] - r[0]
    h2 = h * h
    u = np.zeros_like(r)

    if p.potential_name == "hydrogenic":
        u[0], up0 = hydrogenic_initial_u_up(r[0], p.Z, E, p.l)
    else:
        raise NotImplementedError("一般ポテンシャルの原点条件は個別実装してください。")

    g0 = g_from_Veff(r[0], E, p, Veff_func)
    u[1] = u[0] + h * up0 + 0.5 * h2 * g0 * u[0]

    for i in range(1, len(r) - 1):
        gi = g_from_Veff(r[i], E, p, Veff_func)
        u[i + 1] = 2.0 * u[i] - u[i - 1] + h2 * gi * u[i]
    return u


def integrate_right_verlet(r: np.ndarray, E: float, p: PotentialParams, Veff_func: Callable) -> np.ndarray:
    h = r[1] - r[0]
    h2 = h * h
    u = np.zeros_like(r)

    u[-1], up1 = right_tail_initial_u_up(r[-1], E, p, Veff_func)
    g1 = g_from_Veff(r[-1], E, p, Veff_func)
    u[-2] = u[-1] - h * up1 + 0.5 * h2 * g1 * u[-1]

    for i in range(len(r) - 2, 0, -1):
        gi = g_from_Veff(r[i], E, p, Veff_func)
        u[i - 1] = 2.0 * u[i] - u[i + 1] + h2 * gi * u[i]
    return u


def integrate_left_numerov(r: np.ndarray, E: float, p: PotentialParams, Veff_func: Callable) -> np.ndarray:
    h = r[1] - r[0]
    h2 = h * h
    u = np.zeros_like(r)

    if p.potential_name == "hydrogenic":
        u[0], up0 = hydrogenic_initial_u_up(r[0], p.Z, E, p.l)
    else:
        raise NotImplementedError("一般ポテンシャルの原点条件は個別実装してください。")

    g0 = g_from_Veff(r[0], E, p, Veff_func)
    u[1] = u[0] + h * up0 + 0.5 * h2 * g0 * u[0]

    for i in range(1, len(r) - 1):
        gm = g_from_Veff(r[i - 1], E, p, Veff_func)
        g0 = g_from_Veff(r[i], E, p, Veff_func)
        gp = g_from_Veff(r[i + 1], E, p, Veff_func)
        a = 1.0 - h2 * gp / 12.0
        b = 2.0 * (1.0 + 5.0 * h2 * g0 / 12.0) * u[i]
        c = (1.0 - h2 * gm / 12.0) * u[i - 1]
        u[i + 1] = (b - c) / a

        print(f"E={E:.4g}: i={i} left: u={u[i+1]}")
    return u


def integrate_right_numerov(r: np.ndarray, E: float, p: PotentialParams, Veff_func: Callable) -> np.ndarray:
    h = r[1] - r[0]
    h2 = h * h
    u = np.zeros_like(r)

    u[-1], up1 = right_tail_initial_u_up(r[-1], E, p, Veff_func)
    g1 = g_from_Veff(r[-1], E, p, Veff_func)
    u[-2] = u[-1] - h * up1 + 0.5 * h2 * g1 * u[-1]

    for i in range(len(r) - 2, 0, -1):
        gp = g_from_Veff(r[i + 1], E, p, Veff_func)
        g0 = g_from_Veff(r[i], E, p, Veff_func)
        gm = g_from_Veff(r[i - 1], E, p, Veff_func)
        a = 1.0 - h2 * gm / 12.0
        b = 2.0 * (1.0 + 5.0 * h2 * g0 / 12.0) * u[i]
        c = (1.0 - h2 * gp / 12.0) * u[i + 1]
        u[i - 1] = (b - c) / a

        print(f"E={E:.4g}: i={i} right: u={u[i+1]}")
    return u


LEFT_INTEGRATORS = {
    "euler": integrate_left_euler,
    "verlet": integrate_left_verlet,
    "numerov": integrate_left_numerov,
}
RIGHT_INTEGRATORS = {
    "euler": integrate_right_euler,
    "verlet": integrate_right_verlet,
    "numerov": integrate_right_numerov,
}


# ============================================================
# 接続・解析補助
# ============================================================

def derivative_cdiff(y: np.ndarray, h: float, i: int) -> float:
    if i <= 0:
        return (y[1] - y[0]) / h
    if i >= len(y) - 1:
        return (y[-1] - y[-2]) / h
    return (y[i + 1] - y[i - 1]) / (2.0 * h)


def log_derivative(y: np.ndarray, h: float, i: int, eps: float = 1e-300) -> float:
    return derivative_cdiff(y, h, i) / (y[i] + math.copysign(eps, y[i] if y[i] != 0 else 1.0))


def match_solutions(r: np.ndarray, uL: np.ndarray, uR: np.ndarray, im: int) -> Tuple[np.ndarray, float]:
    """
    中点 im で uL, uR を振幅合わせして1本の u に結合。
    mismatch は対数微分差。
    """
    h = r[1] - r[0]

    if abs(uR[im]) < 1e-300:
        scale = 1.0
    else:
        scale = uL[im] / uR[im]

    uR_scaled = scale * uR

    mismatch = log_derivative(uL, h, im) - log_derivative(uR_scaled, h, im)

    u = np.empty_like(uL)
    u[:im + 1] = uL[:im + 1]
    u[im + 1:] = uR_scaled[im + 1:]
    return u, mismatch


def count_nodes(u: np.ndarray) -> int:
    s = np.sign(u.copy())
    scale = max(np.max(np.abs(u)), 1.0)
    tiny = 1e-14 * scale
    for i in range(len(s)):
        if abs(u[i]) < tiny:
            s[i] = s[i - 1] if i > 0 else 1.0
    nodes = 0
    for i in range(1, len(s)):
        if s[i] * s[i - 1] < 0:
            nodes += 1
    return nodes


def normalize_u(r: np.ndarray, u: np.ndarray) -> np.ndarray:
    norm = np.trapz(u * u, r)
    if norm <= 0 or not np.isfinite(norm):
        return u
    return u / math.sqrt(norm)


def radial_R(r: np.ndarray, u: np.ndarray) -> np.ndarray:
    R = np.zeros_like(u)
    mask = r > 0
    R[mask] = u[mask] / r[mask]
    return R


# ============================================================
# エネルギー探索
# ============================================================

def exact_energy_hydrogenic(Z: float, n: int) -> float:
    return -Z * Z / (2.0 * n * n)


def mismatch_function(
    E: float,
    r: np.ndarray,
    p: PotentialParams,
    Veff_func: Callable,
    method: str,
    im: int,
) -> Tuple[float, np.ndarray]:
    uL = LEFT_INTEGRATORS[method](r[:im + 1], E, p, Veff_func)
    uR = RIGHT_INTEGRATORS[method](r[im:], E, p, Veff_func)

    full_u = np.zeros_like(r)
    matched_u_local, mismatch = match_solutions(r[im:], np.concatenate([uL[-1:]]), uR, 0)  # not used
    # 正しく全体を作る
    if abs(uR[0]) < 1e-300:
        scale = 1.0
    else:
        scale = uL[-1] / uR[0]
    uR_scaled = scale * uR

    h = r[1] - r[0]
    ldL = derivative_cdiff(uL, h, len(uL) - 1) / (uL[-1] + math.copysign(1e-300, uL[-1] if uL[-1] != 0 else 1.0))
    ldR = derivative_cdiff(uR_scaled, h, 0) / (uR_scaled[0] + math.copysign(1e-300, uR_scaled[0] if uR_scaled[0] != 0 else 1.0))
    mismatch = ldL - ldR

    full_u[:im + 1] = uL
    full_u[im:] = uR_scaled
    return mismatch, full_u


def find_bracket_for_mismatch(
    E_center: float,
    r: np.ndarray,
    p: PotentialParams,
    Veff_func: Callable,
    method: str,
    im: int,
    rel_width: float,
    scan_points: int,
) -> Tuple[float, float]:
    width = max(rel_width * abs(E_center), 1e-4)
    Emin = E_center - width
    Emax = min(E_center + width, -1e-10)

    Es = np.linspace(Emin, Emax, scan_points)
    vals = []

    for E in Es:
        try:
            mis, _ = mismatch_function(E, r, p, Veff_func, method, im)
            if np.isfinite(mis):
                vals.append(mis)
            else:
                vals.append(np.nan)
        except Exception:
            vals.append(np.nan)

    # まず通常の符号反転を探す
    for i in range(len(Es) - 1):
        a, b = vals[i], vals[i + 1]
        if np.isfinite(a) and abs(a) < 1e-12:
            return Es[i], Es[i]
        if np.isfinite(a) and np.isfinite(b) and a * b < 0:
            return Es[i], Es[i + 1]

    # 符号反転がない場合、|mismatch|最小点の近傍を返す
    finite_idx = [i for i, v in enumerate(vals) if np.isfinite(v)]
    if not finite_idx:
        raise RuntimeError("mismatch が全点で非数でした。")

    i0 = min(finite_idx, key=lambda i: abs(vals[i]))
    i1 = max(0, i0 - 1)
    i2 = min(len(Es) - 1, i0 + 1)

    if i1 == i2:
        return Es[i0], Es[i0]
    return Es[i1], Es[i2]


def bisect_mismatch(
    E1: float,
    E2: float,
    r: np.ndarray,
    p: PotentialParams,
    Veff_func: Callable,
    method: str,
    im: int,
    tol: float,
    maxiter: int,
) -> float:
    f1, _ = mismatch_function(E1, r, p, Veff_func, method, im)
    f2, _ = mismatch_function(E2, r, p, Veff_func, method, im)

    if abs(f1) < 1e-14:
        return E1
    if abs(f2) < 1e-14:
        return E2
    if f1 * f2 > 0:
        raise RuntimeError("二分法の初期区間で mismatch の符号が変わっていません。")

    a, b = E1, E2
    fa, fb = f1, f2

    for _ in range(maxiter):
        c = 0.5 * (a + b)
        fc, _ = mismatch_function(c, r, p, Veff_func, method, im)

        if abs(fc) < 1e-14 or abs(b - a) < tol:
            return c

        if fa * fc < 0:
            b, fb = c, fc
        else:
            a, fa = c, fc

    return 0.5 * (a + b)


# ============================================================
# 高水準ソルバ
# ============================================================

def solve_bound_state(
    p: PotentialParams,
    s: SolveParams,
    n: Optional[int] = None,
) -> SolveResult:
    Veff_func = get_potential_function(p.potential_name)
    r = np.linspace(s.rmin, s.rmax, s.nmesh)
    im = max(2, min(s.nmesh - 3, int(s.match_fraction * (s.nmesh - 1))))

    if s.solve_e:
        if p.potential_name == "hydrogenic" and n is not None:
            E_center = s.E_guess if s.E_guess is not None else exact_energy_hydrogenic(p.Z, n)
        else:
            if s.E_guess is None:
                raise ValueError("一般ポテンシャルでは E_guess を与えてください。")
            E_center = s.E_guess

        e1, e2 = find_bracket_for_mismatch(
            E_center, r, p, Veff_func, s.method, im,
            s.scan_rel_width, s.scan_points
        )
        E = bisect_mismatch(
            e1, e2, r, p, Veff_func, s.method, im,
            s.bisect_tol, s.bisect_maxiter
        )
    else:
        if s.E_guess is None:
            if p.potential_name == "hydrogenic" and n is not None:
                E = exact_energy_hydrogenic(p.Z, n)
            else:
                raise ValueError("solve_e=False なら E_guess が必要です。")
        else:
            E = s.E_guess

    mismatch, u = mismatch_function(E, r, p, Veff_func, s.method, im)
    u = normalize_u(r, u)
    R = radial_R(r, u)
    nodes = count_nodes(u)

    return SolveResult(
        method=s.method,
        E=E,
        r=r,
        u=u,
        R=R,
        nodes=nodes,
        match_index=im,
        mismatch=mismatch,
    )


# ============================================================
# 水素様原子の解析解
# ============================================================

def assoc_laguerre_integer(k: int, alpha: int, x: np.ndarray) -> np.ndarray:
    """
    整数 k, alpha に対する陪ラゲール多項式 L_k^(alpha)(x)
    再帰で計算
    """
    if k == 0:
        return np.ones_like(x)
    if k == 1:
        return -x + alpha + 1

    Lm2 = np.ones_like(x)
    Lm1 = -x + alpha + 1
    for n in range(2, k + 1):
        Ln = ((2 * n - 1 + alpha - x) * Lm1 - (n - 1 + alpha) * Lm2) / n
        Lm2, Lm1 = Lm1, Ln
    return Lm1


def hydrogenic_radial_exact(r: np.ndarray, Z: float, n: int, l: int) -> np.ndarray:
    """
    解析的な R_{nl}(r)（原子単位系）
    規格化込み。
    """
    if not (0 <= l <= n - 1):
        raise ValueError("0 <= l <= n-1 required")

    rho = 2.0 * Z * r / n
    k = n - l - 1
    alpha = 2 * l + 1

    # 規格化定数
    num = (2.0 * Z / n) ** 3 * math.factorial(k)
    den = 2.0 * n * math.factorial(n + l)
    N = math.sqrt(num / den)

    L = assoc_laguerre_integer(k, alpha, rho)
    R = N * np.exp(-rho / 2.0) * (rho ** l) * L
    return R


def hydrogenic_u_exact(r: np.ndarray, Z: float, n: int, l: int) -> np.ndarray:
    return r * hydrogenic_radial_exact(r, Z, n, l)


# ============================================================
# 可視化
# ============================================================

def plot_solution(
    res_list: List[SolveResult],
    p: PotentialParams,
    n: Optional[int],
    plot_R: bool = False,
    show_exact: bool = True,
) -> None:
    plt.figure(figsize=(8, 5))

    for res in res_list:
        y = res.R if plot_R else res.u
        plt.plot(res.r, y, label=f"{res.method}: E={res.E:.10f}, nodes={res.nodes}")

    if show_exact and p.potential_name == "hydrogenic" and n is not None:
        if plot_R:
            y_exact = hydrogenic_radial_exact(res_list[0].r, p.Z, n, p.l)
        else:
            y_exact = hydrogenic_u_exact(res_list[0].r, p.Z, n, p.l)
        y_exact = normalize_u(res_list[0].r, y_exact if not plot_R else res_list[0].r * y_exact)
        if plot_R:
            y_exact = radial_R(res_list[0].r, y_exact)
        plt.plot(res_list[0].r, y_exact, "--", label="exact")

    plt.xlabel("r [a.u.]")
    plt.ylabel("R(r)" if plot_R else "u(r)=rR(r)")
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.show()


# ============================================================
# main
# ============================================================

def main() -> None:
    ap = argparse.ArgumentParser(description="Practical radial Schrödinger solver with midpoint matching")
    ap.add_argument("--potential", type=str, default="hydrogenic", choices=["hydrogenic"])
    ap.add_argument("--Z", type=float, default=1.0)
    ap.add_argument("--n", type=int, default=1)
    ap.add_argument("--l", type=int, default=0)
    ap.add_argument("--method", type=str, default="numerov", choices=["euler", "verlet", "numerov", "all"])
    ap.add_argument("--solve-e", action="store_true")
    ap.add_argument("--E-guess", type=float, default=None)
    ap.add_argument("--rmin", type=float, default=1e-5)
    ap.add_argument("--rmax", type=float, default=50.0)
    ap.add_argument("--nmesh", type=int, default=6000)
    ap.add_argument("--match-fraction", type=float, default=0.35)
    ap.add_argument("--scan-rel-width", type=float, default=0.4)
    ap.add_argument("--scan-points", type=int, default=500)
    ap.add_argument("--plot", action="store_true")
    ap.add_argument("--plot-R", action="store_true")
    args = ap.parse_args()

    p = PotentialParams(
        potential_name=args.potential,
        Z=args.Z,
        l=args.l,
    )

    methods = ["euler", "verlet", "numerov"] if args.method == "all" else [args.method]
    results: List[SolveResult] = []

    print("=== Radial Schrödinger solver ===")
    print(f"potential = {args.potential}")
    print(f"Z={args.Z}, n={args.n}, l={args.l}")
    if args.potential == "hydrogenic":
        Eex = exact_energy_hydrogenic(args.Z, args.n)
        print(f"exact energy = {Eex:.12f} a.u.")
        print(f"expected nodes = {args.n - args.l - 1}")
    print()

    for m in methods:
        s = SolveParams(
            rmin=args.rmin,
            rmax=args.rmax,
            nmesh=args.nmesh,
            match_fraction=args.match_fraction,
            method=m,
            solve_e=args.solve_e,
            E_guess=args.E_guess,
            scan_rel_width=args.scan_rel_width,
            scan_points=args.scan_points,
        )
        res = solve_bound_state(p, s, n=args.n)
        results.append(res)

        print(f"[{m}]")
        print(f"  E          = {res.E:.12f}")
        if args.potential == "hydrogenic":
            print(f"  abs error  = {abs(res.E - Eex):.6e}")
        print(f"  nodes      = {res.nodes}")
        print(f"  mismatch   = {res.mismatch:.6e}")
        print(f"  match idx  = {res.match_index}")
        print()

    if args.plot or args.plot_R:
        plot_solution(results, p, n=args.n, plot_R=args.plot_R, show_exact=True)


if __name__ == "__main__":
    main()