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

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm, LinearSegmentedColormap
from matplotlib import rcParams
from analyze_res import read_k_resolved

# ===== 全局样式：字体/字号 =====

rcParams["font.sans-serif"] = ["Arial", "DejaVu Sans", "Liberation Sans"]
#rcParams["axes.titlesize"] = 18
#rcParams["axes.labelsize"] = 16
rcParams["xtick.labelsize"] = 20
rcParams["ytick.labelsize"] = 20
#rcParams["legend.fontsize"] = 14

# ========= 配置 =========
path_out  = "k_resolved.out"  # .out 文件
dim1_idx  = 1                   # 0=x, 1=y, 2=z
dim2_idx  = 0
param_idx = 0
k3_idx    = 0
scale     = 1e11                # 仅用于展示/导出（不改变“未放大”的物理总和）

# 实空间晶格（列向量 a1,a2,a3；单位 Å）
a1 = np.array([ 3.618463754, -0.000000008,  0.0])
a2 = np.array([-1.809232280,   3.133681301,  0.0])
a3 = np.array([ 0.0,          0.0,        31.061029434])

EXPORT_DAT = True
png_name = f"heatmap_node_phys_fold_colormap_d1{dim1_idx}_d2{dim2_idx}_k3-{k3_idx}_p-{param_idx}.png"
dat_name = f"nodes_phys_k_fold_d1{dim1_idx}_d2{dim2_idx}_k3-{k3_idx}_p-{param_idx}.dat"
DPI_SAVE  = 300  # 导出分辨率
# =====================================

def cross(u, v):
    return np.array([u[1]*v[2] - u[2]*v[1],
                     u[2]*v[0] - u[0]*v[2],
                     u[0]*v[1] - u[1]*v[0]], dtype=float)

# 1) 倒易基矢（含 2π）与体积
vol = float(np.dot(a1, cross(a2, a3)))      # 实空间体积 V [Å^3]（可能带符号）
b1  = 2*np.pi * cross(a2, a3) / vol         # Å^{-1}
b2  = 2*np.pi * cross(a3, a1) / vol
b3  = 2*np.pi * cross(a1, a2) / vol
BZ_vol = (2*np.pi)**3 / abs(vol)            # |BZ| = (2π)^3/|V|

# 2) 读取 & 选分量
X = read_k_resolved(path_out)  # dims: ('k1','k2','k3','dim1','dim2','param')
for d in ("k1","k2","k3","dim1","dim2","param"):
    if d not in X.dims:
        raise RuntimeError(f"缺少维度 {d}，实际维度: {X.dims}")

A = X.sel(k3=k3_idx, dim1=dim1_idx, dim2=dim2_idx, param=param_idx).values  # (nk1, nk2)
nk1, nk2 = A.shape
nk3 = int(X.sizes["k3"]) if "k3" in X.sizes else 1
Nk  = nk1 * nk2 * nk3

# 3) 半区间折回（与数据同步）：把右半边(>=0.5)移动到负半轴
shift1 = - (nk1 // 2)
shift2 = - (nk2 // 2)
A_fold = np.roll(A, shift=(shift1, shift2), axis=(0, 1))

# 折回后的有序 reduced 结点坐标（与 A_fold 对齐；区间 ~ [-0.5, 0.5) 以步长 1/nk）
f1_axis = (np.arange(nk1) + shift1) / nk1
f2_axis = (np.arange(nk2) + shift2) / nk2

# 固定的 k3（不折回）
f3_nodes = np.arange(nk3, dtype=float) / nk3 if nk3 > 0 else np.array([0.0])
f3_fixed = f3_nodes[k3_idx]

# 物理结点坐标（Å^{-1}），与 A_fold 对应
F1n, F2n = np.meshgrid(f1_axis, f2_axis, indexing="ij")  # (nk1, nk2)
Kn = (F1n[..., None] * b1[None, None, :]
    + F2n[..., None] * b2[None, None, :]
    + f3_fixed        * b3[None, None, :])
kx_nodes = Kn[..., 0]
ky_nodes = Kn[..., 1]
kz_nodes = Kn[..., 2]

# 4) “对所有 k 点取平均 × BZ 体积因子” —— 最终总和（与可视化折回无关）
avg_unscaled   = float(np.nanmean(A))
total_unscaled = avg_unscaled * BZ_vol
A_scaled       = A_fold * scale
avg_scaled     = float(np.nanmean(A_scaled))
total_scaled   = avg_scaled * BZ_vol

print("====== BZ 归一化（平均 × |BZ|），结点热图 + 同步折回 + 自定义色图 ======")
print(f"nk1,nk2,nk3 = {nk1},{nk2},{nk3}  |  Nk = {Nk}")
print(f"V  = {vol:.12e}  [Å^3]   |  |BZ| = (2π)^3/|V| = {BZ_vol:.12e}  [Å^-3]")
print(f"未放大平均 <A>                   = {avg_unscaled:.18e}")
print(f"未放大总和 <A>×|BZ|              = {total_unscaled:.18e}   <-- 最终物理总和")
print(f"放大 {scale:g} 倍后平均 <A×{scale:g}>     = {avg_scaled:.18e}")
print(f"放大后总和 <A×{scale:g}>×|BZ|    = {total_scaled:.18e}")
print("（注意：折回只影响可视化坐标与 A 的排列，不影响上述总和）")
print("======================================")

# 5) 自定义色图：最小=蓝色、0=米黄色、最大=红色
cmap_beige = LinearSegmentedColormap.from_list(
    "blue_beige_red",
    ["#0000FF", "#F5E6C8", "#FF0000"],  # 蓝 → 米黄 → 红
    N=256
)

# 对称色条：以数据最大绝对值为上限，0 居中
clim = float(np.nanmax(np.abs(A_scaled))) if np.size(A_scaled) else 1.0
if not np.isfinite(clim) or clim == 0.0:
    clim = 1.0
norm = TwoSlopeNorm(vmin=-clim, vcenter=0.0, vmax=clim)

# 6) 画图
fig = plt.figure(figsize=(8.6, 7.4))
ax = plt.gca()

mesh = ax.pcolormesh(
    kx_nodes, ky_nodes, A_scaled,
    shading="nearest",
    norm=norm,
    cmap=cmap_beige
)

# 色条
cbar = plt.colorbar(mesh, ax=ax)
cbar.set_label(f"tensor[{dim1_idx},{dim2_idx}] × {scale:g}  (a.u.)")
cbar.set_ticks([-clim, 0.0, clim])
cbar.ax.tick_params(direction="in", length=4, width=1, labelsize=14)
cbar.outline.set_linewidth(1.0)

# 轴样式：刻度线向内、含次刻度、四边都有刻度
ax.tick_params(direction="in", which="both", top=True, right=True, length=6, width=1.1)
#iax.minorticks_on()
#ax.tick_params(which="minor", direction="in", top=True, right=True, length=3, width=1.0)

ax.set_xlabel(r"$k_x$  [$\mathrm{\AA}^{-1}$]")
ax.set_ylabel(r"$k_y$  [$\mathrm{\AA}^{-1}$]")
#ax.set_title(f"Physical-k Node Heatmap (half-folded)  (k3={k3_idx}, param={param_idx})")
ax.set_aspect("equal", adjustable="box")
plt.tight_layout()

# 导出 300 dpi
plt.savefig(png_name, dpi=DPI_SAVE)
print(f"✅ 结点热图已保存: {png_name}  (dpi={DPI_SAVE})")

# 7) 可选导出结点 (kx, ky, kz, value_scaled)
if EXPORT_DAT:
    out = np.column_stack([kx_nodes.ravel(), ky_nodes.ravel(), kz_nodes.ravel(), A_scaled.ravel()])
    np.savetxt(dat_name, out, fmt=["%.10e","%.10e","%.10e","%.18e"])
    print(f"✅ 已导出结点数据: {dat_name}  (行数={out.shape[0]})")
