#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
alpha2eps_iter.py

吸収係数 alpha(E) [cm^-1] から、自己無撞着な KK 変換に基づき
複素誘電関数 eps1(E), eps2(E) と屈折率 n(E), 消衰係数 k(E) を求めるスクリプト。

入力: テキストファイル, 2列: E[eV], alpha[cm^-1]
出力: CSV, 5列: E_eV, n, k, eps1, eps2
"""

import numpy as np
import argparse

# h*c = 1239.841984 eV*nm
HC_EV_NM = 1239.841984


def alpha_to_k(E_eV, alpha_cm):
    """
    E_eV      : array of photon energies [eV]
    alpha_cm  : absorption coefficient [cm^-1]
    return    : k(E) (dimensionless)
    """
    E_eV = np.asarray(E_eV, dtype=float)
    alpha_cm = np.asarray(alpha_cm, dtype=float)

    # lambda [nm] and [cm]
    lambda_nm = HC_EV_NM / E_eV
    lambda_cm = lambda_nm * 1e-7

    k = alpha_cm * lambda_cm / (4.0 * np.pi)
    return k


def kk_eps2_to_eps1(E, eps2, eps_inf=1.0):
    """
    離散 KK 変換:
    eps1(E) = eps_inf + (2/pi) * P∫ [E' eps2(E') / (E'^2 - E^2)] dE'

    E      : energy grid [eV], 1D array, 昇順を想定
    eps2   : imaginary part eps2(E) on the same grid
    eps_inf: high-energy limit of eps1

    return : eps1(E)
    """
    E = np.asarray(E, dtype=float)
    eps2 = np.asarray(eps2, dtype=float)

    # ensure ascending order
    idx = np.argsort(E)
    E = E[idx]
    eps2 = eps2[idx]

    # effective grid spacing for numerical integration
    dE = np.empty_like(E)
    dE[1:-1] = 0.5 * (E[2:] - E[:-2])
    dE[0] = E[1] - E[0]
    dE[-1] = E[-1] - E[-2]

    eps1 = np.empty_like(eps2)

    for j, Ej in enumerate(E):
        denom = E**2 - Ej**2
        # principal value: skip singularity by zeroing denom==0
        with np.errstate(divide="ignore", invalid="ignore"):
            kernel = np.where(
                denom != 0.0,
                E * eps2 / denom,
                0.0
            )
        integral = np.sum(kernel * dE)
        eps1[j] = eps_inf + (2.0 / np.pi) * integral

    return eps1, idx


def solve_self_consistent(E, alpha_cm, n_init=1.5, eps_inf=1.0,
                          max_iter=50, tol=1e-4, verbose=True):
    """
    自己無撞着 KK で n(E) を求める。

    Parameters
    ----------
    E        : array, energies [eV]
    alpha_cm : array, absorption coefficient [cm^-1]
    n_init   : initial guess of n(E) (constant)
    eps_inf  : high-energy limit of eps1
    max_iter : maximum number of iterations
    tol      : convergence criterion for max|n_new - n_old|
    verbose  : print iteration info

    Returns
    -------
    E_sorted : sorted energy array
    n        : converged n(E)
    k        : k(E) from alpha
    eps1     : eps1(E)
    eps2     : eps2(E)
    """
    E = np.asarray(E, dtype=float)
    alpha_cm = np.asarray(alpha_cm, dtype=float)

    # sort by energy
    idx = np.argsort(E)
    E = E[idx]
    alpha_cm = alpha_cm[idx]

    # k from alpha (fixed in all iterations)
    k = alpha_to_k(E, alpha_cm)

    # initial n(E)
    n = np.full_like(E, float(n_init))

    for it in range(max_iter):
        # eps2 from current n and fixed k
        eps2 = 2.0 * n * k

        # KK to get eps1
        eps1_tmp, idx_kk = kk_eps2_to_eps1(E, eps2, eps_inf=eps_inf)

        # kk_eps2_to_eps1 内で再ソートしているので、元の順序に戻す
        eps1 = np.empty_like(eps1_tmp)
        eps1[idx_kk] = eps1_tmp

        # update n from eps1 and k:
        # n_new(E) = sqrt(eps1(E) + k(E)^2)
        inside = eps1 + k**2
        # 数値誤差で inside が負になるのを防ぐ
        inside = np.maximum(inside, 0.0)
        n_new = np.sqrt(inside)

        diff = np.max(np.abs(n_new - n))
        if verbose:
            print(f"Iteration {it+1:3d}: max|Δn| = {diff:.4e}")

        n = n_new

        if diff < tol:
            if verbose:
                print("Converged.")
            break
    else:
        if verbose:
            print("WARNING: did not converge within max_iter.")

    # final eps2, eps1 from KK (self-consistent pair)
    eps2 = 2.0 * n * k
    eps1_tmp, idx_kk = kk_eps2_to_eps1(E, eps2, eps_inf=eps_inf)
    eps1 = np.empty_like(eps1_tmp)
    eps1[idx_kk] = eps1_tmp

    return E, n, k, eps1, eps2


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Absorption coefficient alpha(E) -> self-consistent complex "
            "dielectric function eps1(E), eps2(E) via KK."
        )
    )
    parser.add_argument("infile", help="Input text file: E[eV]  alpha[cm^-1]")
    parser.add_argument("-o", "--outfile", default="eps_self_consistent.csv",
                        help="Output CSV file: E_eV, n, k, eps1, eps2")
    parser.add_argument("--n_init", type=float, default=1.5,
                        help="Initial constant refractive index n(E)")
    parser.add_argument("--eps_inf", type=float, default=1.0,
                        help="High-energy limit of eps1(E)")
    parser.add_argument("--max_iter", type=int, default=50,
                        help="Maximum number of iterations")
    parser.add_argument("--tol", type=float, default=1e-4,
                        help="Convergence criterion for max|Δn|")
    parser.add_argument("--quiet", action="store_true",
                        help="Suppress iteration output")

    args = parser.parse_args()

    data = np.loadtxt(args.infile)
    if data.shape[1] < 2:
        raise ValueError("Input file must have at least 2 columns: E[eV] alpha[cm^-1]")

    E = data[:, 0]
    alpha = data[:, 1]

    E_out, n, k, eps1, eps2 = solve_self_consistent(
        E,
        alpha,
        n_init=args.n_init,
        eps_inf=args.eps_inf,
        max_iter=args.max_iter,
        tol=args.tol,
        verbose=not args.quiet,
    )

    out = np.column_stack([E_out, n, k, eps1, eps2])
    header = "E_eV,n,k,eps1,eps2"
    np.savetxt(args.outfile, out, delimiter=",", header=header, comments="")

    print(f"Saved self-consistent dielectric function to {args.outfile}")


if __name__ == "__main__":
    main()