"""
数据读取模块 - 从 POSCAR 文件读取结构信息和参数
"""

import os
from typing import Optional, Tuple, List, Dict
import numpy as np
import warnings


def read_poscar(filename="POSCAR", to_direct: bool=True, use_scale: bool=True):
    """
    读取扩展的 POSCAR 文件格式，支持从第一行读取 dimk, dimr, hopping_decay 参数
    
    新的 POSCAR 格式：
    第1行: dimr=3 dimk=3 [hopping_decay=1.0] 或 3 3 1 (dimk dimr hopping_decay)
    第2行: scale factor
    第3-5行: 晶格向量
    第6行: 元素列表
    第7行: 元素数量
    第8行: Direct/Cartesian
    第9行开始: 原子坐标 (x y z [magmom])
    
    参数:
        filename (str): POSCAR文件路径，默认为"POSCAR"
        to_direct (bool): 是否将坐标转换为分数坐标
        use_scale (bool): 是否使用比例因子
        
    返回:
        dict: 包含结构信息和参数的字典
    """
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            if not lines:
                raise ValueError(f"文件 {filename} 为空")

        # 解析第1行：dimr, dimk, hopping_decay
        first_line = lines[0].strip()
        dimk, dimr, hopping_decay = _parse_first_line(first_line)
        
        # 解析第2行：scale factor
        try:
            scale_factor = float(lines[1].strip())
        except (ValueError, IndexError):
            raise ValueError(f"POSCAR 文件 {filename} 的第二行比例因子无效: '{lines[1].strip() if len(lines) > 1 else 'N/A'}'")
        scale = scale_factor if use_scale else 1.0

        # 解析第3-5行：晶格向量
        if len(lines) < 5:
            raise ValueError(f"POSCAR 文件 {filename} 行数不足，无法读取晶格向量")
        lattice = np.array([list(map(float, line.split())) for line in lines[2:5]]) * scale
        
        # 计算最小晶格常数（用于 max_distance）
        lattice_lengths = np.linalg.norm(lattice, axis=1)
        min_lattice_constant = float(np.min(lattice_lengths))
        
        # 解析第6行：元素列表
        if len(lines) < 6:
            raise ValueError(f"POSCAR 文件 {filename} 缺少元素列表行")
        elements_line = lines[5].strip()
        elements = elements_line.split()
        
        # 解析第7行：元素数量
        if len(lines) < 7:
            raise ValueError(f"POSCAR 文件 {filename} 缺少元素数量行")
        counts_line = lines[6].strip()
        try:
            atom_counts = list(map(int, counts_line.split()))
        except ValueError:
            raise ValueError(f"POSCAR 文件 {filename} 的第七行原子数无效: '{counts_line}'")

        if len(elements) != len(atom_counts):
            raise ValueError(f"POSCAR 文件 {filename} 中元素名数量 ({len(elements)}) 与原子数数量 ({len(atom_counts)}) 不匹配。")

        total_atoms = sum(atom_counts)
        
        # 解析第8行：坐标类型
        if len(lines) < 8:
            raise ValueError(f"POSCAR 文件 {filename} 缺少坐标类型行 (第 8 行)")
        coord_type = lines[7].strip()
        is_direct = False
        if coord_type.lower().startswith('d'):
            is_direct = True
        elif coord_type.lower().startswith('c') or coord_type.lower().startswith('k'):
            is_direct = False
        elif coord_type.lower().startswith('s'):
            # Selective dynamics
            if len(lines) < 9:
                raise ValueError(f"POSCAR 文件 {filename} 在 Selective dynamics 后缺少坐标类型行")
            coord_type = lines[8].strip()
            if coord_type.lower().startswith('d'):
                is_direct = True
            elif coord_type.lower().startswith('c') or coord_type.lower().startswith('k'):
                is_direct = False
            else:
                raise ValueError(f"未知的坐标类型 (在 Selective dynamics 之后): {coord_type}")
        else:
            warnings.warn(f"POSCAR 文件 {filename} 坐标类型行 (第 8 行) 未识别: '{coord_type}'. 假设为 Direct 坐标。")
            is_direct = True

        coord_start_line = 8 if not coord_type.lower().startswith('s') else 9

        # 读取原子坐标和磁矩
        if coord_start_line + total_atoms > len(lines):
            raise ValueError(f"POSCAR 文件 {filename} 行数不足 ({len(lines)})，无法读取 {total_atoms} 个原子坐标 (从第 {coord_start_line+1} 行开始)")
        
        coords_raw = []
        magnetic_moments = []
        for i in range(total_atoms):
            try:
                parts = lines[coord_start_line + i].split()
                coords_raw.append(list(map(float, parts[:3])))
                # 读取磁矩（如果存在第4列）
                if len(parts) >= 4:
                    try:
                        magmom = float(parts[3])
                        magnetic_moments.append(magmom)
                    except ValueError:
                        magnetic_moments.append(0.0)
                else:
                    magnetic_moments.append(0.0)
            except (ValueError, IndexError):
                raise ValueError(f"POSCAR 文件 {filename} 的第 {coord_start_line + i + 1} 行坐标格式错误: '{lines[coord_start_line + i].strip()}'")
        
        coords = np.array(coords_raw)
        magnetic_moments = np.array(magnetic_moments)

        # 生成每个原子的元素符号列表
        atom_symbols_all = []
        for symbol, count in zip(elements, atom_counts):
            atom_symbols_all.extend([symbol] * count)

        # 坐标转换
        if not is_direct and to_direct:
            try:
                inv_lattice = np.linalg.inv(lattice)
                coords = np.dot(coords, inv_lattice)
                is_direct = True
            except np.linalg.LinAlgError:
                raise ValueError(f"POSCAR 文件 {filename} 的晶格向量无法求逆，无法将笛卡尔坐标转换为分数坐标。")
        elif is_direct and not to_direct:
            coords = np.dot(coords, lattice)
            is_direct = False

        # 构建结果字典
        poscar_data = {
            "dimk": dimk,
            "dimr": dimr,
            "hopping_decay": hopping_decay,
            "min_lattice_constant": min_lattice_constant,
            "scale": scale_factor,
            "lattice": lattice / scale_factor if use_scale else lattice,
            "elements": elements,
            "atom_counts": atom_counts,
            "total_atoms": total_atoms,
            "atom_symbols": atom_symbols_all,
            "coordinates": coords,
            "magnetic_moments": magnetic_moments,
            "is_direct": is_direct,
        }
        
        # 打印读取信息
        import sys
        if hasattr(sys.stdout, 'write'):
            print(f"成功读取 POSCAR 文件: {filename}")
            print(f"  元素: {elements}")
            print(f"  原子数: {atom_counts} (总计: {total_atoms})")
            print(f"  晶格常数: {min_lattice_constant:.4f} 埃")
            print(f"  参数: dimk={dimk}, dimr={dimr}, hopping_decay={hopping_decay}")
            if np.any(np.abs(magnetic_moments) > 1e-10):
                print(f"  磁矩: 已读取 {np.sum(np.abs(magnetic_moments) > 1e-10)} 个非零磁矩")
        
        return poscar_data
    except FileNotFoundError:
        raise FileNotFoundError(f"POSCAR 文件 {filename} 未找到。")
    except Exception as e:
        print(f"读取POSCAR文件 '{filename}' 时出错: {str(e)}")
        raise


def _parse_first_line(first_line: str) -> Tuple[int, int, float]:
    """
    解析 POSCAR 第一行，提取 dimk, dimr, hopping_decay
    
    支持两种格式：
    1. "dimr=3 dimk=3 hopping_decay=1.0"
    2. "3 3 1" (dimk dimr hopping_decay)
    
    返回:
        (dimk, dimr, hopping_decay)
    """
    # 尝试解析 key=value 格式
    if '=' in first_line:
        dimk = 3  # 默认值
        dimr = 3  # 默认值
        hopping_decay = 1.0  # 默认值
        
        parts = first_line.split()
        for part in parts:
            if '=' in part:
                key, value = part.split('=', 1)
                key = key.strip().lower()
                try:
                    if key == 'dimk':
                        dimk = int(value)
                    elif key == 'dimr':
                        dimr = int(value)
                    elif key == 'hopping_decay':
                        hopping_decay = float(value)
                except ValueError:
                    warnings.warn(f"无法解析参数 '{part}'，使用默认值")
        
        return dimk, dimr, hopping_decay
    else:
        # 尝试解析空格分隔的数值格式
        try:
            parts = first_line.split()
            if len(parts) >= 2:
                dimk = int(parts[0])
                dimr = int(parts[1])
                hopping_decay = float(parts[2]) if len(parts) >= 3 else 1.0
                return dimk, dimr, hopping_decay
            else:
                raise ValueError(f"第一行格式错误，需要至少2个数值: '{first_line}'")
        except ValueError:
            warnings.warn(f"无法解析第一行 '{first_line}'，使用默认值: dimk=3, dimr=3, hopping_decay=1.0")
            return 3, 3, 1.0


def read_kpath(filename: Optional[str] = None) -> Tuple[List[List[float]], List[str], int]:
    """
    从 KPATH.in 或 KPOINTS 文件读取 k 路径
    
    参数:
        filename (str, optional): k路径文件路径。如果为 None，则尝试查找 KPATH.in 或 KPOINTS
        
    返回:
        tuple: (kpath, klabels, num_k_points)
            - kpath: k点坐标列表
            - klabels: k点标签列表
            - num_k_points: k点数量（从文件第2行读取，如果不存在则返回默认值100）
            
    异常:
        FileNotFoundError: 如果找不到 k 路径文件
    """
    # 如果没有指定文件名，尝试查找
    if filename is None:
        if os.path.exists("KPATH.in"):
            filename = "KPATH.in"
        elif os.path.exists("KPOINTS"):
            filename = "KPOINTS"
        else:
            raise FileNotFoundError("找不到 k 路径文件。请提供 KPATH.in 或 KPOINTS 文件。")
    
    if not os.path.exists(filename):
        raise FileNotFoundError(f"k 路径文件 '{filename}' 未找到。")
    
    kpath = []
    klabels = []
    num_k_points = 100  # 默认值
    
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        if len(lines) < 4:
            raise ValueError(f"KPATH 文件 '{filename}' 行数过少，无法解析路径。")
        
        # 读取第2行的 k 点数量（如果存在且是数字）
        if len(lines) >= 2:
            try:
                second_line = lines[1].strip()
                # 尝试解析为整数
                num_k_points = int(second_line)
            except (ValueError, IndexError):
                # 如果第2行不是数字，使用默认值
                warnings.warn(f"KPATH 文件 '{filename}' 的第2行不是有效的 k 点数量，使用默认值 100")
                num_k_points = 100
        
        # 查找坐标行开始位置（跳过头部）
        coord_line_start = 0
        for i, line in enumerate(lines):
            parts = line.strip().split()
            if len(parts) >= 3:
                try:
                    # 检查前3个部分是否为浮点数
                    float(parts[0])
                    float(parts[1])
                    float(parts[2])
                    coord_line_start = i
                    break
                except ValueError:
                    continue
            if i > 10:
                raise ValueError(f"无法在 KPATH 文件 '{filename}' 中找到坐标行。")

        if coord_line_start == 0 and not any(len(line.strip().split()) >= 3 for line in lines):
            raise ValueError(f"未能在 KPATH 文件 '{filename}' 中找到坐标行。")

        # 处理坐标行
        path_segments = []
        current_segment = []
        last_kpoint = None

        for i in range(coord_line_start, len(lines)):
            line = lines[i].strip()
            
            # 空行作为段分隔符
            if not line:
                if current_segment:
                    path_segments.append(current_segment)
                    current_segment = []
                continue

            parts = line.split()
            if len(parts) < 3:
                warnings.warn(f"KPATH 文件 '{filename}' 第 {i+1} 行格式无效: '{line}'")
                continue

            try:
                coords = [float(x) for x in parts[:3]]
            except ValueError:
                warnings.warn(f"KPATH 文件 '{filename}' 第 {i+1} 行坐标无效: '{line}'")
                continue

            # 提取标签（通常在 '!' 后或第4个元素）
            label = ""
            if '!' in line:
                label = line.split('!', 1)[1].strip()
            elif len(parts) > 3:
                try:
                    float(parts[3])  # 如果是数字，不是标签
                except ValueError:
                    label = parts[3]

            current_segment.append({"coords": coords, "label": label})
            last_kpoint = {"coords": coords, "label": label}

        # 添加最后一段
        if current_segment:
            path_segments.append(current_segment)

        if not path_segments:
            raise ValueError(f"未能从 KPATH 文件 '{filename}' 中解析出任何路径段。")

        # 组合路径段
        first_point = path_segments[0][0]
        kpath.append(first_point["coords"])
        klabels.append(first_point["label"])

        for segment in path_segments:
            if len(segment) > 1:
                end_point = segment[-1]
                if not kpath or not np.allclose(kpath[-1], end_point["coords"]):
                    kpath.append(end_point["coords"])
                    klabels.append(end_point["label"])
                elif klabels[-1] == "":
                    klabels[-1] = end_point["label"]

        if len(kpath) != len(klabels):
            warnings.warn(f"KPATH 文件 '{filename}' 解析后 k 点数与标签数不匹配。")

        return kpath, klabels, num_k_points
        
    except FileNotFoundError:
        raise
    except Exception as e:
        raise ValueError(f"读取 KPATH 文件 '{filename}' 时发生错误: {str(e)}")

