"""
参数管理类 - 从 POSCAR 文件读取参数
"""

import numpy as np
from typing import List, Optional
from .read_datas import read_poscar, read_kpath
import warnings


class Parameters:
    """
    参数管理类，从 POSCAR 文件读取所有参数
    """
    
    def __init__(self, poscar_filename: str = "POSCAR"):
        """
        从 POSCAR 文件初始化参数
        
        参数:
            poscar_filename (str): POSCAR 文件路径
        """
        # 读取 POSCAR 文件
        self.poscar_data = read_poscar(poscar_filename)
        
        # 从 POSCAR 数据初始化参数
        self._initialize_from_poscar()
        
        # 读取 k 路径
        self._read_kpath()
        
        # 验证参数
        self._validate_parameters()
    
    def _initialize_from_poscar(self):
        """从 POSCAR 数据初始化所有参数"""
        data = self.poscar_data
        
        # 从 POSCAR 第一行读取的参数
        self.dimk = int(data["dimk"])
        self.dimr = int(data["dimr"])
        self.hopping_decay = float(data["hopping_decay"])
        
        # 晶格相关
        self.a0 = 1.0  # 晶格常数缩放因子（POSCAR 中已应用 scale）
        self.lattice = data["lattice"]
        self.min_lattice_constant = data["min_lattice_constant"]
        
        # 距离参数（固定值）
        self.min_distance = 0.1
        self.max_distance = self.min_lattice_constant  # 最大距离 = 最大晶格常数
        self.max_R_range = 1  # 最大相邻格点搜索范围
        
        # 跃迁参数（默认值，可以从 POSCAR 读取）
        self.t0 = -1.0  # 默认参考跃迁强度
        self.t0_distance = 2.5  # 默认参考距离
        
        # 自旋和磁性
        self.nspin = 2  # 默认有自旋
        self.sigma_z = np.array([[1, 0], [0, -1]])
        self.magnetic_moments = data["magnetic_moments"]  # 从 POSCAR 读取的磁矩
        
        # 元素和坐标
        self.elements = data["elements"]
        self.atom_counts = data["atom_counts"]
        self.total_atoms = data["total_atoms"]
        self.atom_symbols = data["atom_symbols"]
        self.coordinates = data["coordinates"]
        
        # 其他参数
        self.same_atom_negative_coupling = False
        self.onsite_energy = np.zeros(self.total_atoms)  # 默认在位能为0
        
        # 输出参数
        self.output_filename = "band_structure"
        self.output_format = "png"
        self.savedir = "."
        
        # 绘图参数
        self.ylim = [-2, 2]
        self.is_black_degenerate_bands = True
        self.energy_threshold = 1e-4
        
        # 输出控制
        self.is_print_tb_model_hop = False
        self.is_print_tb_model = True
        self.is_check_flat_bands = False
        self.is_report_kpath = False
    
    def _read_kpath(self):
        """读取 k 路径文件"""
        try:
            self.kpath, self.klabel, num_k_points = read_kpath()
            self.num_k_points = num_k_points  # 从文件读取的 k 点数量
            import sys
            if hasattr(sys.stdout, 'write'):
                print(f"成功读取 k 路径文件")
                print(f"  k 路径点数: {len(self.kpath)}")
                print(f"  路径标签: {self.klabel}")
                print(f"  k 点数量（用于插值）: {self.num_k_points}")
        except FileNotFoundError as e:
            raise FileNotFoundError(f"无法读取 k 路径文件: {e}")
    
    def _validate_parameters(self):
        """验证参数有效性"""
        # 验证 dimk 和 dimr
        if self.dimk not in [1, 2, 3]:
            raise ValueError(f"dimk 必须是 1, 2 或 3，当前值: {self.dimk}")
        if self.dimr not in [1, 2, 3]:
            raise ValueError(f"dimr 必须是 1, 2 或 3，当前值: {self.dimr}")
        
        # 验证 k 路径维度
        if self.kpath:
            kpath_array = np.array(self.kpath)
            if kpath_array.shape[1] > self.dimk:
                # 截断到 dimk 维度
                self.kpath = kpath_array[:, :self.dimk].tolist()
            elif kpath_array.shape[1] < self.dimk:
                # 用0填充
                warnings.warn(f"k 路径维度 ({kpath_array.shape[1]}) 小于 dimk ({self.dimk})，将用0填充")
                padded_kpath = np.pad(kpath_array, 
                                    ((0, 0), (0, self.dimk - kpath_array.shape[1])),
                                    mode='constant',
                                    constant_values=0)
                self.kpath = padded_kpath.tolist()
        
        # 验证磁矩长度
        if len(self.magnetic_moments) != self.total_atoms:
            warnings.warn(f"磁矩数量 ({len(self.magnetic_moments)}) 与原子数 ({self.total_atoms}) 不匹配")
    
    def get_maglist(self) -> List[float]:
        """
        获取磁矩列表
        
        返回:
            List[float]: 磁矩列表
        """
        return self.magnetic_moments.tolist()
    
    def __repr__(self):
        """返回参数的字符串表示"""
        return f"""Parameters(
    POSCAR: {self.poscar_data.get('scale', 'N/A')}
    dimk={self.dimk}, dimr={self.dimr}
    hopping_decay={self.hopping_decay}
    max_distance={self.max_distance:.4f} (max lattice constant)
    max_R_range={self.max_R_range}
    total_atoms={self.total_atoms}
    elements={self.elements}
)"""

