"""
紧束缚模型计算模块 - 利用 pythtb 库建立紧束缚模型并计算能带结构
"""

from pythtb import *  # import TB model class
import numpy as np
import matplotlib.pyplot as plt
import os
from .parameters import Parameters
from copy import deepcopy


def hopping_strength(distance, params):
    """
    计算跃迁强度与距离的关系
    
    参数:
        distance (float): 实际距离（埃）
        params (Parameters): 参数实例
        
    返回:
        float: 跃迁强度
    """
    # 使用指数衰减模型: t = t0 * exp(-lambda_*(d-d0)/d0)
    return params.t0 * np.exp(-params.hopping_decay * (distance - params.t0_distance) / params.t0_distance)


def get_coupling_strength(atom1_index, atom2_index, poscar_data, params):
    """
    计算两个原子之间的耦合强度，考虑周期性边界条件和相邻格点上的等价原子
    
    参数:
        atom1_index (int): 第一个原子的索引（0-based）
        atom2_index (int): 第二个原子在数组中的索引（0-based）
        poscar_data (dict): POSCAR数据字典
        params (Parameters): 参数实例
        
    返回:
        tuple: (耦合强度数组, 格矢量数组, 距离数组)
    """
    # 获取原子坐标
    coord1 = poscar_data["coordinates"][atom1_index]
    coord2 = poscar_data["coordinates"][atom2_index]

    # 获取元素类型
    type1 = poscar_data["atom_symbols"][atom1_index]
    type2 = poscar_data["atom_symbols"][atom2_index]

    # 如果两个原子是同一种元素，则耦合强度为正，否则为负（可选）
    coupling_sign = -1 if ((type1 == type2) and params.same_atom_negative_coupling) else 1
    
    # 生成相邻格点的格矢量
    R_vectors = []
    coupling_values = []
    distance_values = []

    # 生成所有可能的格矢量
    Rlist = []
    if params.dimk == 3:
        for i in range(-params.max_R_range, params.max_R_range+1):
            for j in range(-params.max_R_range, params.max_R_range+1):
                for k in range(-params.max_R_range, params.max_R_range+1):
                    Rlist.append(np.array([i, j, k]))
    elif params.dimk == 2:
        for i in range(-params.max_R_range, params.max_R_range+1):
            for j in range(-params.max_R_range, params.max_R_range+1):
                Rlist.append(np.array([i, j, 0]))
    elif params.dimk == 1:
        Rlist = [np.array([i, 0, 0]) for i in range(-params.max_R_range, params.max_R_range+1)]
    else:
        raise ValueError(f"dimk 必须是 1, 2 或 3，当前值: {params.dimk}")
    
    # 考虑周期性边界条件下的相邻格点
    for R in Rlist:
        # 计算考虑周期性的距离（分数坐标）
        diff = coord2 + R - coord1
        
        # 转换为笛卡尔坐标
        cart_diff = np.dot(diff, poscar_data["lattice"])
        
        # 计算笛卡尔坐标下的距离
        distance = np.linalg.norm(cart_diff)

        if (distance > params.max_distance) or (distance < params.min_distance):
            # 超过最大距离或小于最小距离则跳过
            continue
        
        # 计算耦合强度
        coupling = hopping_strength(distance, params)
        
        # 存储结果
        coupling_values.append(coupling_sign * coupling)
        distance_values.append(distance)
        R_vectors.append(R)
    
    return coupling_values, R_vectors, distance_values


def calculate_all_couplings(poscar_data, params):
    """
    计算所有原子对之间的跃迁强度和向量
    
    参数:
        poscar_data (dict): POSCAR数据字典
        params (Parameters): 参数实例
        
    返回:
        list: 包含所有耦合信息的列表
    """
    n_atoms = len(poscar_data["coordinates"])
    all_couplings = []
    
    # 计算所有可能的原子对之间的耦合
    for i in range(n_atoms):
        for j in range(i, n_atoms):
            # 计算耦合强度和格矢量
            coupling_values, R_vectors, distance_values = get_coupling_strength(
                i, j, poscar_data, params
            )
            
            # 存储结果
            coupling_info = {
                "atom1_index": i,
                "atom2_index": j,
                "coupling_values": coupling_values,
                "distance_values": distance_values,
                "R_vectors": R_vectors
            }
            
            all_couplings.append(coupling_info)
    
    return all_couplings


def remove_duplicate_hoppings(couplings):
    """
    消除重复的跃迁对
    
    参数:
        couplings (list): 跃迁信息列表
        
    返回:
        list: 去除重复后的跃迁信息列表
    """
    unique_hoppings = []
    seen_hoppings = set()
    
    for hop in couplings:
        for t, R, d in zip(hop['coupling_values'], hop['R_vectors'], hop['distance_values']):
            # 创建标准化的跃迁标识符
            if hop['atom1_index'] <= hop['atom2_index']:
                atom_pair = (hop['atom1_index'], hop['atom2_index'])
                R_vec = tuple(R)
            else:
                atom_pair = (hop['atom2_index'], hop['atom1_index'])
                R_vec = tuple(-r for r in R)
                
            hopping_id = (atom_pair, R_vec)
            neg_hopping_id = (atom_pair, tuple(-r for r in R_vec))
            
            # 如果这个跃迁或其反向跃迁都没见过，则添加
            if hopping_id not in seen_hoppings and neg_hopping_id not in seen_hoppings:
                seen_hoppings.add(hopping_id)
                new_hop = {
                    "atom1_index": hop['atom1_index'],
                    "atom2_index": hop['atom2_index'],
                    "coupling_values": [t],
                    "distance_values": [d],
                    "R_vectors": [list(R)]
                }
                unique_hoppings.append(new_hop)
    return unique_hoppings


def create_pythtb_model(params: Parameters):
    """
    根据 POSCAR 数据创建 pythtb 模型
    
    参数:
        params (Parameters): 参数实例

    返回:
        pythtb.tb_model: 创建的紧束缚模型
    """
    try:
        from pythtb import tb_model
    except ImportError:
        raise ImportError("请安装 pythtb 库: pip install pythtb==1.8.0")
    
    # 使用 POSCAR 数据
    poscar_data = params.poscar_data
    
    # 提取晶格向量
    lattice = poscar_data['lattice'] / params.a0
    
    # 提取原子坐标（已经是分数坐标）
    coords = poscar_data['coordinates']
    
    # 使用距离基方法创建紧束缚模型
    model = tb_model(params.dimk, params.dimr, lattice, coords, nspin=params.nspin)
    
    # 计算所有原子对之间的耦合
    print(f"计算原子对之间的耦合（共 {len(coords)} 个原子）...")
    all_couplings = calculate_all_couplings(poscar_data, params)
    print(f"找到 {len(all_couplings)} 个原子对")
    
    all_couplings = remove_duplicate_hoppings(all_couplings)
    print(f"去除重复后剩余 {len(all_couplings)} 个跃迁")
    
    # 统计跃迁数量
    total_hoppings = sum(len(hop['coupling_values']) for hop in all_couplings)
    print(f"总跃迁数量: {total_hoppings}")
    print()
    
    # 获取原子数量
    num_atoms_used = len(coords)
    
    # 设置在位能（初始化为0）
    for ind in range(num_atoms_used):
        model.set_onsite(0, ind)
    
    # 设置磁性（从 POSCAR 读取的磁矩）
    maglist = params.get_maglist()
    for ind in range(min(len(maglist), num_atoms_used)):
        if abs(maglist[ind]) > 1e-10:  # 只设置非零磁矩
            model.set_onsite(maglist[ind] * params.sigma_z, ind, "add")
    
    # 设置跃迁参数
    for hop in all_couplings:
        for t, R, d in zip(hop['coupling_values'], hop['R_vectors'], hop['distance_values']):
            model.set_hop(t, hop['atom1_index'], hop['atom2_index'], 
                         [int(R[0]), int(R[1]), int(R[2])], 
                         allow_conjugate_pair=True)
            if params.is_print_tb_model_hop:
                print(f"model.set_hop({t:.6f}, {hop['atom1_index']}, {hop['atom2_index']}, "
                      f"[{R[0]}, {R[1]}, {R[2]}]) # distance: {d:.4f}")

    if params.is_print_tb_model:
        model.display()
    
    return model


def adjust_degenerate_bands(evals, evecs, model, energy_threshold=1e-3):
    """
    修改简并能带的自旋投影。如果找到简并的能带且自旋投影相反，则将其投影设为0。
    
    参数:
        evals (numpy.ndarray): 能带本征值
        evecs (numpy.ndarray): 能带本征矢量
        model (pythtb.tb_model): 紧束缚模型
        energy_threshold (float): 判断能带简并的能量阈值
        
    返回:
        tuple: (evals, modified_evecs)
    """
    n_bands, n_kpoints = evals.shape
    modified_evecs = deepcopy(evecs)
    
    # 遍历每个k点
    for k in range(n_kpoints):
        # 遍历每个能带
        for band1 in range(n_bands):
            for band2 in range(band1 + 1, n_bands):
                # 检查能带是否简并
                if abs(evals[band1, k] - evals[band2, k]) < energy_threshold:
                    # 计算两个能带的自旋投影
                    spin_up1, spin_down1 = 0.0, 0.0
                    spin_up2, spin_down2 = 0.0, 0.0
                    
                    for orb in range(model._norb):
                        spin_up1 += abs(evecs[band1, k, orb, 0])**2
                        spin_down1 += abs(evecs[band1, k, orb, 1])**2
                        spin_up2 += abs(evecs[band2, k, orb, 0])**2
                        spin_down2 += abs(evecs[band2, k, orb, 1])**2
                    
                    total1 = spin_up1 + spin_down1
                    total2 = spin_up2 + spin_down2
                    
                    if total1 > 1e-10 and total2 > 1e-10:
                        spin_up1 /= total1
                        spin_down1 /= total1
                        spin_up2 /= total2
                        spin_down2 /= total2
                        
                        # 检查自旋投影是否相反
                        if (abs(spin_up1 - spin_down2) < 0.001 and 
                            abs(spin_down1 - spin_up2) < 0.001):
                            # 将两个能带的自旋投影都设为0
                            for orb in range(model._norb):
                                modified_evecs[band1, k, orb] = 0
                                modified_evecs[band2, k, orb] = 0
    
    return evals, modified_evecs


def calculate_band_structure(model, params: Parameters):
    """
    计算能带结构并绘图
    
    参数:
        model (pythtb.tb_model): 紧束缚模型
        params (Parameters): 参数实例
    """
    # 计算能带
    (k_vec, k_dist, k_node) = model.k_path(params.kpath, params.num_k_points, 
                                            report=params.is_report_kpath)
    evals, evecs = model.solve_all(k_vec, eig_vectors=True)
    
    # 调整简并能带
    if params.is_black_degenerate_bands:
        evals, evecs = adjust_degenerate_bands(evals, evecs, model, params.energy_threshold)
    
    # 计算自旋投影
    spin_up_projection = np.zeros_like(evals)
    spin_down_projection = np.zeros_like(evals)

    # 对于每个能带和每个k点
    for band in range(evals.shape[0]):
        for k in range(evals.shape[1]):
            # 计算自旋上投影
            up_weight = 0.0
            down_weight = 0.0
            for orb in range(model._norb):
                up_weight += np.abs(evecs[band, k, orb, 0])**2
                down_weight += np.abs(evecs[band, k, orb, 1])**2
                
            # 归一化
            total = up_weight + down_weight
            if total > 1e-10:
                spin_up_projection[band, k] = up_weight / total
                spin_down_projection[band, k] = down_weight / total
    
    # 绘制带有自旋颜色的能带
    fig, ax = plt.subplots(figsize=(10, 6))

    # 设置x轴
    ax.set_xlim(k_node[0], k_node[-1])
    ax.set_xticks(k_node)
    ax.set_xticklabels(params.klabel)
    for n in range(len(k_node)):
        ax.axvline(x=k_node[n], linewidth=0.5, color='k')

    # 用颜色渐变绘制能带
    for band in range(evals.shape[0]):
        # 创建一个颜色映射: 红色=自旋上, 蓝色=自旋下
        colors = []
        for k in range(evals.shape[1]):
            r = spin_up_projection[band, k]
            b = spin_down_projection[band, k]
            colors.append([r, 0, b])
        
        # 使用渐变色绘制线段
        for k in range(evals.shape[1]-1):
            ax.plot(k_dist[k:k+2], evals[band, k:k+2], 
                    c=colors[k], linewidth=2)
    
    # 设置x轴刻度
    ax.set_xticks(k_node)
    ax.set_xticklabels(params.klabel)
    
    # 添加垂直线
    for node in k_node:
        ax.axvline(x=node, linewidth=0.5, color='k')
    
    # 设置标签和范围
    ax.set_xlabel('k-path')
    ax.set_ylabel('Energy (eV)')
    ax.set_ylim(params.ylim)
    
    # 保存图片
    output_path = os.path.join(params.savedir, f"{params.output_filename}.{params.output_format}")
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"能带结构图已保存到: {output_path}")
    print(f"能带数量: {evals.shape[0]}")
    print(f"k 点数量: {evals.shape[1]}")
    print(f"能量范围: [{np.min(evals):.4f}, {np.max(evals):.4f}] eV")

