#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
anharmonic_chain_Tshift.py

1D monatomic chain (nearest-neighbor) anharmonic phonon frequency shift vs temperature.

Model (relative displacement Δ_n = u_{n+1} - u_n):
  V = Σ_n [ (K2/2) Δ_n^2 + (K3/3!) Δ_n^3 + (K4/4!) Δ_n^4 ]

We compute (perturbatively):
  ω_q(T) ≈ ω_q + Δω_q^(4)(T) + Δω_q^(3)(T)
where
  Δω_q ≈ Re Σ(q, ω_q) / (2 ω_q)

- Quartic (K4) "tadpole"-like shift: simple thermal average of ⟨|u_k|^2⟩.
- Cubic (K3) "bubble" shift: principal value sum over k, regularized by small eta.

Units:
- Use SI by default (a in meters, M in kg, K2 N/m, K3 N/m^2, K4 N/m^3).
- Frequencies are in rad/s internally; output can be converted to THz if desired.

You can customize K2/K3/K4 as functions of (delta, T) by editing k2_func/k3_func/k4_func.
"""

from __future__ import annotations
import argparse
import math
from dataclasses import dataclass
from typing import Callable, Tuple, List

import numpy as np

# Physical constants (SI)
HBAR = 1.054_571_817e-34  # J*s
KB   = 1.380_649e-23      # J/K
TWOPI = 2.0 * math.pi


# -----------------------------
# User-customizable force constants (as functions)
# -----------------------------
def k2_func(delta: float = 0.0, T: float | None = None) -> float:
    """
    Harmonic force constant K2 [N/m] (or consistent unit if you use dimensionless convention).
    Replace with any function you like.
    """
    return 10.0  # example


def k3_func(delta: float = 0.0, T: float | None = None) -> float:
    """
    Cubic anharmonic constant K3 [N/m^2].
    Replace with any function you like.
    """
    return 0.0  # example (set nonzero to enable bubble shift)


def k4_func(delta: float = 0.0, T: float | None = None) -> float:
    """
    Quartic anharmonic constant K4 [N/m^3].
    Replace with any function you like.
    """
    return 1.0e-5  # example (set nonzero to enable tadpole shift)


# -----------------------------
# Core model helpers
# -----------------------------
def C_of_q(q: np.ndarray, a: float) -> np.ndarray:
    """
    C(q) = e^{iqa} - 1
    We'll mostly use |C(q)|^2 = 4 sin^2(qa/2), which is real.
    """
    return np.exp(1j * q * a) - 1.0


def absC2(q: np.ndarray, a: float) -> np.ndarray:
    """|C(q)|^2 = 4 sin^2(qa/2)."""
    return 4.0 * np.sin(0.5 * q * a) ** 2


def omega_harmonic(q: np.ndarray, a: float, M: float, K2: float) -> np.ndarray:
    """
    Harmonic dispersion (nearest neighbor):
      ω_q^2 = (4 K2 / M) sin^2(qa/2)
    """
    return np.sqrt((4.0 * K2 / M) * np.sin(0.5 * q * a) ** 2)


def bose_n(omega: np.ndarray, T: float, use_quantum: bool = True) -> np.ndarray:
    """
    Bose factor n(ω) = 1/(exp(β ħ ω) - 1).
    If use_quantum=False, use classical limit: n ~ kT/(ħ ω).
    """
    if T <= 0.0:
        return np.zeros_like(omega)
    if not use_quantum:
        # classical limit
        return (KB * T) / (HBAR * np.maximum(omega, 1e-300))
    x = (HBAR * omega) / (KB * T)
    # avoid overflow for large x
    x = np.clip(x, 0.0, 700.0)
    return 1.0 / (np.expm1(x) + 1e-300)


def coth_halfbeta_hw(omega: np.ndarray, T: float, use_quantum: bool = True) -> np.ndarray:
    """
    coth(βħω/2). If T->0, -> 1.
    Classical limit: coth ~ 2kT/(ħω).
    """
    if T <= 0.0:
        return np.ones_like(omega)
    if not use_quantum:
        return (2.0 * KB * T) / (HBAR * np.maximum(omega, 1e-300))
    x = (HBAR * omega) / (2.0 * KB * T)
    x = np.clip(x, 0.0, 350.0)
    # coth(x) = cosh(x)/sinh(x) = (e^x + e^-x)/(e^x - e^-x)
    ex = np.exp(x)
    emx = 1.0 / ex
    return (ex + emx) / (ex - emx + 1e-300)


# -----------------------------
# Quartic (K4) tadpole-like shift
# -----------------------------
def delta_omega_quartic_tadpole(
    q: float,
    kgrid: np.ndarray,
    a: float,
    M: float,
    K2: float,
    K4: float,
    T: float,
    use_quantum: bool,
    k_cut: float,
) -> float:
    """
    A practical 1D implementation of the structure:
      Σ^(4)(q) ∝ K4 |C(q)|^2 * (1/N) Σ_k |C(k)|^2 <|u_k|^2>

    <|u_k|^2> = ħ/(2 M ω_k) coth(βħω_k/2)

    We return Δω_q^(4) = ReΣ/(2ω_q) with a simple prefactor choice.
    NOTE: The exact numerical prefactor depends on convention/diagram counting.
          Here we provide a consistent *model* formula with an adjustable prefactor.
    """
    omega_k = omega_harmonic(kgrid, a, M, K2)
    omega_q = float(omega_harmonic(np.array([q]), a, M, K2)[0])

    if omega_q <= 0.0:
        return 0.0

    mask = np.abs(kgrid) >= k_cut
    kk = kgrid[mask]
    omk = omega_k[mask]

    # <|u_k|^2>
    u2 = (HBAR / (2.0 * M * np.maximum(omk, 1e-300))) * coth_halfbeta_hw(omk, T, use_quantum=use_quantum)

    S = np.mean(absC2(kk, a) * u2)  # (1/N) Σ_k ...
    # prefactor (model choice)
    # Σ^(4) ~ (K4 / 2) |C(q)|^2 * S
    Sigma4 = 0.5 * K4 * absC2(np.array([q]), a)[0] * S
    d_omega = Sigma4 / (2.0 * omega_q)
    return float(d_omega)


# -----------------------------
# Cubic (K3) bubble-like shift
# -----------------------------
def v3_squared(
    q: float,
    k: np.ndarray,
    a: float,
    M: float,
    K3: float,
    omega_q: float,
    omega_k: np.ndarray,
    omega_qmk: np.ndarray,
) -> np.ndarray:
    """
    |V3(q,k,q-k)|^2 up to a prefactor.
    Structure:
      V3 ∝ K3 * C(q)C(k)C(q-k) / sqrt( (2M)^3 ω_q ω_k ω_{q-k} )
    We'll use |C|^2 = 4 sin^2(...) to avoid complex phases.
    """
    # use |C| instead of C itself (phases cancel in |V|^2)
    num = absC2(np.array([q]), a)[0] * absC2(k, a) * absC2((q - k), a)
    den = (2.0 * M) ** 3 * np.maximum(omega_q * omega_k * omega_qmk, 1e-300)
    return (K3 ** 2) * num / den


def re_sigma_cubic_bubble(
    q: float,
    kgrid: np.ndarray,
    a: float,
    M: float,
    K2: float,
    K3: float,
    T: float,
    use_quantum: bool,
    eta: float,
    k_cut: float,
) -> float:
    """
    Compute Re Σ^(3)(q, ω_q) using the standard bubble structure, regularized by eta:
      Re[1/(x + iη)] = x/(x^2 + η^2)

    Σ^(3)(q, ω) = Σ_k |V3|^2 [ (n_k + n_{q-k} + 1)/(ω - ω_k - ω_{q-k} + i0)
                             + (n_k - n_{q-k})/(ω - ω_k + ω_{q-k} + i0)
                             + (n_{q-k} - n_k)/(ω + ω_k - ω_{q-k} + i0)
                             + (n_k + n_{q-k} + 1)/(ω + ω_k + ω_{q-k} + i0) ]

    Return Re Σ at ω=ω_q.
    """
    omega_q = float(omega_harmonic(np.array([q]), a, M, K2)[0])
    if omega_q <= 0.0:
        return 0.0

    mask = np.abs(kgrid) >= k_cut
    k = kgrid[mask]
    omega_k = omega_harmonic(k, a, M, K2)
    omega_qmk = omega_harmonic(q - k, a, M, K2)

    nk = bose_n(omega_k, T, use_quantum=use_quantum)
    nqm = bose_n(omega_qmk, T, use_quantum=use_quantum)

    V2 = v3_squared(q, k, a, M, K3, omega_q, omega_k, omega_qmk)

    def re_inv(x: np.ndarray) -> np.ndarray:
        return x / (x * x + eta * eta)

    w = omega_q

    x1 = w - omega_k - omega_qmk
    x2 = w - omega_k + omega_qmk
    x3 = w + omega_k - omega_qmk
    x4 = w + omega_k + omega_qmk

    term1 = (nk + nqm + 1.0) * re_inv(x1)
    term2 = (nk - nqm)       * re_inv(x2)
    term3 = (nqm - nk)       * re_inv(x3)
    term4 = (nk + nqm + 1.0) * re_inv(x4)

    Sigma3 = np.sum(V2 * (term1 + term2 + term3 + term4)) / k.size
    return float(Sigma3)


def delta_omega_cubic_bubble(
    q: float,
    kgrid: np.ndarray,
    a: float,
    M: float,
    K2: float,
    K3: float,
    T: float,
    use_quantum: bool,
    eta: float,
    k_cut: float,
) -> float:
    """
    Δω_q^(3) = Re Σ^(3)(q, ω_q) / (2 ω_q)
    """
    omega_q = float(omega_harmonic(np.array([q]), a, M, K2)[0])
    if omega_q <= 0.0:
        return 0.0
    reS = re_sigma_cubic_bubble(q, kgrid, a, M, K2, K3, T, use_quantum, eta, k_cut)
    return reS / (2.0 * omega_q)


# -----------------------------
# Main calculation
# -----------------------------
@dataclass
class Params:
    a: float
    M: float
    Nk: int
    k_cut: float
    use_quantum: bool
    eta_rel: float


def make_kgrid(Nk: int, a: float) -> np.ndarray:
    """
    Uniform k grid in the first Brillouin zone: (-pi/a, pi/a]
    """
    k = np.linspace(-math.pi / a, math.pi / a, Nk, endpoint=False)
    return k


def omega_T_for_q(
    q: float,
    T: float,
    kgrid: np.ndarray,
    p: Params,
    k2f: Callable[[float, float | None], float],
    k3f: Callable[[float, float | None], float],
    k4f: Callable[[float, float | None], float],
) -> Tuple[float, float, float, float]:
    """
    Return (omega0, domega3, domega4, omegaT) for given q and T.
    """
    K2 = float(k2f(0.0, T))
    K3 = float(k3f(0.0, T))
    K4 = float(k4f(0.0, T))

    omega0 = float(omega_harmonic(np.array([q]), p.a, p.M, K2)[0])

    # regularization scale for principal value
    eta = max(p.eta_rel * omega0, 1e-12)

    d3 = 0.0
    d4 = 0.0
    if abs(K3) > 0.0 and omega0 > 0.0:
        d3 = delta_omega_cubic_bubble(q, kgrid, p.a, p.M, K2, K3, T, p.use_quantum, eta, p.k_cut)

    if abs(K4) > 0.0 and omega0 > 0.0:
        d4 = delta_omega_quartic_tadpole(q, kgrid, p.a, p.M, K2, K4, T, p.use_quantum, p.k_cut)

    omegaT = omega0 + d3 + d4
    return omega0, d3, d4, omegaT


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Anharmonic phonon frequency shift vs T (1D chain, K2/K3/K4 as functions).")
    ap.add_argument("--a", type=float, default=1.0, help="Lattice constant a [m] (or your unit).")
    ap.add_argument("--M", type=float, default=1.0, help="Mass M [kg] (or your unit).")
    ap.add_argument("--Nk", type=int, default=4096, help="Number of k-points in BZ.")
    ap.add_argument("--kcut", type=float, default=0.0, help="Exclude |k| < kcut to avoid IR issues (1D). Unit: 1/a.")
    ap.add_argument("--q", type=float, default=0.5, help="Target q as fraction of pi/a (q = frac * pi/a).")
    ap.add_argument("--Tmin", type=float, default=0.0, help="Minimum temperature [K].")
    ap.add_argument("--Tmax", type=float, default=300.0, help="Maximum temperature [K].")
    ap.add_argument("--nT", type=int, default=31, help="Number of temperature points.")
    ap.add_argument("--classical", action="store_true", help="Use classical (high-T) occupations instead of quantum.")
    ap.add_argument("--eta_rel", type=float, default=1e-3, help="Regularization eta = eta_rel * omega0 for bubble principal value.")
    ap.add_argument("--out", type=str, default="omega_vs_T.csv", help="Output CSV filename.")
    ap.add_argument("--thz", action="store_true", help="Output frequencies in THz (cycles/s) instead of rad/s.")
    return ap.parse_args()


def main() -> None:
    args = parse_args()

    a = args.a
    M = args.M
    Nk = args.Nk

    q = args.q * (math.pi / a)  # q in 1/m (or 1/a)
    k_cut = args.kcut * (1.0 / a) if args.kcut > 0 else 0.0

    p = Params(
        a=a,
        M=M,
        Nk=Nk,
        k_cut=k_cut,
        use_quantum=(not args.classical),
        eta_rel=args.eta_rel,
    )

    kgrid = make_kgrid(Nk, a)

    Ts = np.linspace(args.Tmin, args.Tmax, args.nT)

    rows: List[List[float]] = []
    for T in Ts:
        omega0, d3, d4, omegaT = omega_T_for_q(q, float(T), kgrid, p, k2_func, k3_func, k4_func)
        rows.append([float(T), omega0, d3, d4, omegaT])

    arr = np.array(rows, dtype=float)

    # unit conversion for output
    # rad/s -> THz (cycles/s) = omega/(2π)/1e12
    if args.thz:
        conv = 1.0 / (TWOPI * 1e12)
        arr[:, 1:] *= conv

    header = "T,omega0,domega3_bubble,domega4_tadpole,omegaT"
    np.savetxt(args.out, arr, delimiter=",", header=header, comments="")
    print(f"Wrote: {args.out}")
    print("Columns:", header)
    if args.thz:
        print("Frequencies are in THz.")
    else:
        print("Frequencies are in rad/s.")


if __name__ == "__main__":
    main()