#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Minimal demonstration of the spin magnetization generated by a nonunitary
equal-spin triplet pairing state.

The model is a spin-degenerate square-lattice band with chiral p-wave
equal-spin gaps

    Delta_upup(k,T)     = Delta(T) f(k),
    Delta_downdown(k,T) = r Delta(T) f(k),

where f(k) = sin(kx) + i sin(ky).

r = 1: unitary state.
0 <= r < 1: nonunitary state.
r = 0: maximally nonunitary single-spin-pairing limit.

The electronic magnetization is calculated from the BdG occupations,

    n_{k,sigma} = 1/2 [1 - xi_k/E_{k,sigma}
                       tanh(E_{k,sigma}/2T)],

    M_z = mu_B/N sum_k [n_{k,up} - n_{k,down}].

The chemical potential is adjusted at every temperature to keep the total
particle density fixed.  All energies are measured in units of t.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import json
import numpy as np


@dataclass(frozen=True)
class ModelParams:
    t: float = 1.0
    t2: float = 0.35
    mu_normal: float = 1.0
    Tc: float = 0.25
    Delta0: float = 0.60
    mu_B: float = 1.0


@dataclass(frozen=True)
class NumericalParams:
    Nk: int = 301
    n_temperature: int = 81
    Tmin_ratio: float = 0.02
    Tmax_ratio: float = 1.20
    r_values: tuple[float, ...] = (1.0, 0.6, 0.25, 0.0)
    representative_r: float = 0.25
    representative_T_ratio: float = 0.05
    dos_points: int = 500
    dos_broadening: float = 0.015
    density_temperature: float = 1.0e-4
    root_tolerance: float = 1.0e-12
    root_max_iter: int = 120


@dataclass(frozen=True)
class OutputParams:
    root_dir: str = "nonunitary_magnetization_results"


def make_kmesh(Nk: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    k = np.linspace(-np.pi, np.pi, Nk, endpoint=False)
    kx, ky = np.meshgrid(k, k, indexing="ij")
    return k, kx, ky


def bare_dispersion(
    kx: np.ndarray,
    ky: np.ndarray,
    p: ModelParams,
) -> np.ndarray:
    return (
        -2.0 * p.t * (np.cos(kx) + np.cos(ky))
        - 4.0 * p.t2 * np.cos(kx) * np.cos(ky)
    )


def chiral_p_abs2(kx: np.ndarray, ky: np.ndarray) -> np.ndarray:
    """|sin(kx) + i sin(ky)|^2."""
    return np.sin(kx) ** 2 + np.sin(ky) ** 2


def bcs_gap(T: float, p: ModelParams) -> float:
    """Standard BCS interpolation used only for this pedagogical demo."""
    if T >= p.Tc:
        return 0.0
    if T <= 0.0:
        return p.Delta0
    x = max(p.Tc / T - 1.0, 0.0)
    return p.Delta0 * np.tanh(1.74 * np.sqrt(x))


def fermi(x: np.ndarray, T: float) -> np.ndarray:
    if T <= 1.0e-12:
        return (x < 0.0).astype(float)
    z = np.clip(x / T, -700.0, 700.0)
    return 1.0 / (np.exp(z) + 1.0)


def bdg_occupation(
    xi: np.ndarray,
    gap_abs: np.ndarray,
    T: float,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Return electron occupation n_k and positive BdG energy E_k for one spin
    sector.
    """
    E = np.sqrt(xi * xi + gap_abs * gap_abs)
    ratio = np.divide(xi, E, out=np.zeros_like(xi), where=E > 1.0e-14)
    if T <= 1.0e-12:
        thermal = np.ones_like(E)
    else:
        thermal = np.tanh(E / (2.0 * T))
    n = 0.5 * (1.0 - ratio * thermal)
    return n, E


def total_density(
    mu: float,
    T: float,
    r: float,
    eps0: np.ndarray,
    form_abs: np.ndarray,
    p: ModelParams,
) -> float:
    xi = eps0 - mu
    Delta = bcs_gap(T, p)
    n_up, _ = bdg_occupation(xi, Delta * form_abs, T)
    n_down, _ = bdg_occupation(xi, r * Delta * form_abs, T)
    return float(np.mean(n_up + n_down))


def solve_mu_fixed_density(
    target_density: float,
    T: float,
    r: float,
    eps0: np.ndarray,
    form_abs: np.ndarray,
    p: ModelParams,
    n: NumericalParams,
) -> float:
    lo = float(np.min(eps0) - 4.0 * p.Delta0 - 4.0 * p.Tc)
    hi = float(np.max(eps0) + 4.0 * p.Delta0 + 4.0 * p.Tc)

    f_lo = total_density(lo, T, r, eps0, form_abs, p) - target_density
    f_hi = total_density(hi, T, r, eps0, form_abs, p) - target_density

    if not (f_lo <= 0.0 <= f_hi):
        raise RuntimeError(
            "Chemical-potential bracket does not contain the target density."
        )

    for _ in range(n.root_max_iter):
        mid = 0.5 * (lo + hi)
        f_mid = total_density(mid, T, r, eps0, form_abs, p) - target_density
        if abs(f_mid) < n.root_tolerance:
            return mid
        if f_mid > 0.0:
            hi = mid
        else:
            lo = mid
    return 0.5 * (lo + hi)


def state_observables(
    T: float,
    r: float,
    target_density: float,
    eps0: np.ndarray,
    form_abs: np.ndarray,
    form_abs2: np.ndarray,
    p: ModelParams,
    n: NumericalParams,
) -> dict[str, float | np.ndarray]:
    mu = solve_mu_fixed_density(
        target_density, T, r, eps0, form_abs, p, n
    )
    xi = eps0 - mu
    Delta = bcs_gap(T, p)

    gap_up_abs = Delta * form_abs
    gap_down_abs = r * Delta * form_abs

    n_up, E_up = bdg_occupation(xi, gap_up_abs, T)
    n_down, E_down = bdg_occupation(xi, gap_down_abs, T)

    magnetization = p.mu_B * float(np.mean(n_up - n_down))
    q_gap_map = 0.5 * (
        gap_up_abs * gap_up_abs - gap_down_abs * gap_down_abs
    )
    q_gap_average = float(np.mean(q_gap_map))

    pair_norm = float(
        np.mean(gap_up_abs * gap_up_abs + gap_down_abs * gap_down_abs)
    )
    normalized_pair_spin = (
        2.0 * q_gap_average / pair_norm if pair_norm > 1.0e-15 else 0.0
    )

    return {
        "mu": mu,
        "Delta_up": Delta,
        "Delta_down": r * Delta,
        "density_up": float(np.mean(n_up)),
        "density_down": float(np.mean(n_down)),
        "magnetization": magnetization,
        "q_gap_average": q_gap_average,
        "normalized_pair_spin": normalized_pair_spin,
        "xi": xi,
        "n_up_map": n_up,
        "n_down_map": n_down,
        "delta_n_map": n_up - n_down,
        "q_gap_map": q_gap_map,
        "E_up": E_up,
        "E_down": E_down,
        "gap_up_abs": gap_up_abs,
        "gap_down_abs": gap_down_abs,
    }


def gaussian_dos(
    energies: np.ndarray,
    grid: np.ndarray,
    eta: float,
) -> np.ndarray:
    dos = np.zeros_like(grid)
    flat = energies.ravel()
    chunk = 10000
    norm = 1.0 / (np.sqrt(2.0 * np.pi) * eta * flat.size)
    for start in range(0, flat.size, chunk):
        e = flat[start : start + chunk]
        x = grid[:, None] - e[None, :]
        dos += norm * np.sum(np.exp(-0.5 * (x / eta) ** 2), axis=1)
    return dos


def run_calculation() -> None:
    p = ModelParams()
    n = NumericalParams()
    out = OutputParams()

    root = Path(out.root_dir)
    data_dir = root / "data"
    data_dir.mkdir(parents=True, exist_ok=True)

    k, kx, ky = make_kmesh(n.Nk)
    eps0 = bare_dispersion(kx, ky, p)
    form_abs2 = chiral_p_abs2(kx, ky)
    form_abs = np.sqrt(form_abs2)

    # Fix the total filling to the low-temperature normal-state value.
    normal_xi = eps0 - p.mu_normal
    target_density = float(
        np.mean(
            2.0 * fermi(normal_xi, n.density_temperature)
        )
    )

    temperatures = np.linspace(
        n.Tmin_ratio * p.Tc,
        n.Tmax_ratio * p.Tc,
        n.n_temperature,
    )

    columns = [
        "r",
        "T",
        "T_over_Tc",
        "mu",
        "Delta_up",
        "Delta_down",
        "n_up",
        "n_down",
        "M_z_muB_per_site",
        "q_gap_average",
        "normalized_pair_spin",
    ]
    rows: list[list[float]] = []

    for r in n.r_values:
        for T in temperatures:
            obs = state_observables(
                T, r, target_density, eps0, form_abs, form_abs2, p, n
            )
            rows.append(
                [
                    r,
                    T,
                    T / p.Tc,
                    float(obs["mu"]),
                    float(obs["Delta_up"]),
                    float(obs["Delta_down"]),
                    float(obs["density_up"]),
                    float(obs["density_down"]),
                    float(obs["magnetization"]),
                    float(obs["q_gap_average"]),
                    float(obs["normalized_pair_spin"]),
                ]
            )

    array = np.asarray(rows, dtype=float)
    np.savetxt(
        data_dir / "temperature_scan.dat",
        array,
        header=" ".join(columns),
        fmt="%.10f",
    )
    np.savez_compressed(
        data_dir / "temperature_scan.npz",
        data=array,
        columns=np.array(columns, dtype="U"),
        r_values=np.asarray(n.r_values),
        temperatures=temperatures,
        target_density=target_density,
    )

    # Low-temperature scan versus nonunitarity parameter r.
    r_scan = np.linspace(0.0, 1.0, 81)
    T_low = n.representative_T_ratio * p.Tc
    r_rows = []
    for r in r_scan:
        obs = state_observables(
            T_low, r, target_density, eps0, form_abs, form_abs2, p, n
        )
        r_rows.append(
            [
                r,
                1.0 - r * r,
                float(obs["mu"]),
                float(obs["magnetization"]),
                float(obs["q_gap_average"]),
                float(obs["normalized_pair_spin"]),
            ]
        )
    r_array = np.asarray(r_rows)
    np.savetxt(
        data_dir / "nonunitarity_scan.dat",
        r_array,
        header=(
            "r eta_nonunitary mu M_z_muB_per_site "
            "q_gap_average normalized_pair_spin"
        ),
        fmt="%.10f",
    )

    # Representative momentum maps and spin-resolved quasiparticle DOS.
    r0 = n.representative_r
    T0 = n.representative_T_ratio * p.Tc
    obs0 = state_observables(
        T0, r0, target_density, eps0, form_abs, form_abs2, p, n
    )
    np.savez_compressed(
        data_dir / "representative_maps.npz",
        k=k,
        kx=kx,
        ky=ky,
        xi=np.asarray(obs0["xi"]),
        delta_n_map=np.asarray(obs0["delta_n_map"]),
        q_gap_map=np.asarray(obs0["q_gap_map"]),
        E_up=np.asarray(obs0["E_up"]),
        E_down=np.asarray(obs0["E_down"]),
        gap_up_abs=np.asarray(obs0["gap_up_abs"]),
        gap_down_abs=np.asarray(obs0["gap_down_abs"]),
        r=r0,
        T=T0,
        mu=float(obs0["mu"]),
        magnetization=float(obs0["magnetization"]),
        q_gap_average=float(obs0["q_gap_average"]),
    )

    emax = max(
        1.2 * p.Delta0,
        float(np.percentile(np.asarray(obs0["E_up"]), 12.0)),
        float(np.percentile(np.asarray(obs0["E_down"]), 12.0)),
    )
    energy_grid = np.linspace(0.0, emax, n.dos_points)
    dos_up = gaussian_dos(
        np.asarray(obs0["E_up"]), energy_grid, n.dos_broadening
    )
    dos_down = gaussian_dos(
        np.asarray(obs0["E_down"]), energy_grid, n.dos_broadening
    )
    np.savetxt(
        data_dir / "representative_spin_dos.dat",
        np.column_stack([energy_grid, dos_up, dos_down]),
        header="energy DOS_up DOS_down",
        fmt="%.10f",
    )

    metadata = {
        "model": {
            "t": p.t,
            "t2": p.t2,
            "mu_normal": p.mu_normal,
            "Tc": p.Tc,
            "Delta0": p.Delta0,
            "mu_B": p.mu_B,
        },
        "numerics": {
            "Nk": n.Nk,
            "n_temperature": n.n_temperature,
            "r_values": list(n.r_values),
            "representative_r": n.representative_r,
            "representative_T_over_Tc": n.representative_T_ratio,
            "dos_broadening": n.dos_broadening,
        },
        "target_density": target_density,
        "pairing": {
            "Delta_upup": "Delta(T) [sin(kx)+i sin(ky)]",
            "Delta_downdown": "r Delta(T) [sin(kx)+i sin(ky)]",
            "unitary_condition": "r = 1",
            "nonunitary_condition": "0 <= r < 1",
        },
    }
    (data_dir / "parameters.json").write_text(
        json.dumps(metadata, indent=2), encoding="utf-8"
    )

    # Human-readable summary.
    lines = [
        "Nonunitary equal-spin pairing magnetization demonstration",
        "========================================================",
        "",
        "Pairing:",
        "  Delta_upup(k,T)     = Delta(T) [sin(kx)+i sin(ky)]",
        "  Delta_downdown(k,T) = r Delta(T) [sin(kx)+i sin(ky)]",
        "",
        f"Target density = {target_density:.10f}",
        f"Representative T/Tc = {n.representative_T_ratio:.6f}",
        "",
    ]
    for r in n.r_values:
        obs = state_observables(
            T_low, r, target_density, eps0, form_abs, form_abs2, p, n
        )
        lines.extend(
            [
                f"r = {r:.2f}",
                f"  M_z = {float(obs['magnetization']):.10f} mu_B/site",
                f"  <q_Delta,z> = {float(obs['q_gap_average']):.10f}",
                (
                    "  normalized pair spin = "
                    f"{float(obs['normalized_pair_spin']):.10f}"
                ),
                "",
            ]
        )
    (data_dir / "calculation_summary.txt").write_text(
        "\n".join(lines), encoding="utf-8"
    )

    print(f"Calculation finished. Data saved in: {data_dir}")


if __name__ == "__main__":
    run_calculation()
