import numpy as np
import matplotlib.pyplot as plt
from scipy import interpolate
from matplotlib.path import Path

# -----------------------------
# 1) 读取坐标
# -----------------------------
points = np.loadtxt('PtBi2-kslice-coord.dat')
points_x = np.around(points[:, 0], decimals=10)
points_y = np.around(points[:, 1], decimals=10)
num_pt = len(points)

xmin, xmax = points_x.min(), points_x.max()
ymin, ymax = points_y.min(), points_y.max()

# -----------------------------
# 2) 构造真实范围的规则网格（不强制正方形）
# -----------------------------
num_int = int(round(np.sqrt(num_pt)))

spanx = max(xmax - xmin, 1e-15)
spany = max(ymax - ymin, 1e-15)
span_max = max(spanx, spany)

nx = max(2, int(round(num_int * spanx / span_max)))
ny = max(2, int(round(num_int * spany / span_max)))

xint = np.linspace(xmin, xmax, nx)
yint = np.linspace(ymin, ymax, ny)
grid_x, grid_y = np.meshgrid(xint, yint)

# -----------------------------
# 3) 读取数值
# -----------------------------
val = np.loadtxt('PtBi2-kslice-shc.dat', usecols=(0,)) * 3.45

# -----------------------------
# 4) 数值变换（保持原来的逻辑）
# -----------------------------
val_log = np.array([
    np.log10(abs(v)) * np.sign(v) if abs(v) > 10 else v / 10.0
    for v in val
])

# -----------------------------
# 5) 删除原始数据中 y = 0 的点，再做最近邻插值
# -----------------------------
tol = 1e-10  # 容差，防止浮点数精度问题

mask = np.abs(points_y) > tol

points_x_use = points_x[mask]
points_y_use = points_y[mask]
val_log_use = val_log[mask]

print("原始点数:", len(points_y))
print("去掉 y=0 后剩余点数:", len(points_y_use))
print("被去掉的 y=0 点数:", np.sum(~mask))

valint = interpolate.griddata(
    (points_x_use, points_y_use),
    val_log_use,
    (grid_x, grid_y),
    method='nearest'
)

# -----------------------------
# 6) 定义六边形（你给的 6 个点）
# -----------------------------
hex_points = np.array([
    [ 0.63797375,  0.0       ],
    [ 0.31898688,  0.55269188],
    [-0.31898688,  0.55269188],
    [-0.63797375,  0.0       ],
    [-0.31898688, -0.55269188],
    [ 0.31898688, -0.55269188],
], dtype=float)

hex_path = Path(hex_points)  # 自动视为闭合多边形

# -----------------------------
# 7) 做掩膜：六边形外设为 NaN（透明）
# -----------------------------
xy = np.column_stack([grid_x.ravel(), grid_y.ravel()])
inside = hex_path.contains_points(xy).reshape(grid_x.shape)

val_hex = valint.copy()
val_hex[~inside] = np.nan  # 六边形外变成 NaN

# -----------------------------
# 8) 画“只含六边形内部”的图（PNG）
# -----------------------------
vmax = np.nanmax(np.abs(val_hex))
vmin = -vmax

fig, ax = plt.subplots(figsize=(7, 6))

im = ax.imshow(
    val_hex,
    origin='lower',
    extent=(xmin, xmax, ymin, ymax),
    cmap='bwr',
    vmin=vmin,
    vmax=vmax,
    aspect='equal'
)

cbar = fig.colorbar(im, ax=ax)
cbar.set_label('SHC (scaled)', fontsize=12)

# 画六边形边界
hex_closed = np.vstack([hex_points, hex_points[0]])
ax.plot(
    hex_closed[:, 0], hex_closed[:, 1],
    linewidth=2.0,
    color='k'
)

# 让视野只聚焦六边形附近（留一点边距）
padx = 0.05 * (hex_points[:, 0].max() - hex_points[:, 0].min())
pady = 0.05 * (hex_points[:, 1].max() - hex_points[:, 1].min())

ax.set_xlim(hex_points[:, 0].min() - padx, hex_points[:, 0].max() + padx)
ax.set_ylim(hex_points[:, 1].min() - pady, hex_points[:, 1].max() + pady)

ax.axis('off')

outfile = 'PtBi2-kslice-shc_hex_only_no_y0_interp.png'
plt.savefig(outfile, dpi=300, bbox_inches='tight', transparent=True)
plt.show()
