Rustb 0.7.0 - Docs.rs

#[cfg_attr(doc, katexit::katexit)]
use crate::atom_struct::{Atom, AtomType, OrbProj, OrbitalId};
use crate::error::{Result, TbError};
use crate::{HasRMatrix, Model, RMatrixData, find_R};
use ndarray::prelude::*;
use ndarray_linalg::*;
use num_complex::Complex;

const BOHR_TO_ANGSTROM: f64 = 0.529_177_210_67;

fn win_data_line(line: &str) -> &str {
    // Wannier90 treats both `!` and `#` as inline comment markers.
    // Strip whichever comes first so that commented-out directives
    // (e.g. `#spinors = .true.`) are not mistaken for active ones.
    let line = line.split('!').next().unwrap_or_default();
    line.split('#').next().unwrap_or_default().trim()
}

fn length_unit_scale(token: &str) -> Option<f64> {
    match token.to_ascii_lowercase().as_str() {
        "ang" | "angstrom" | "angstroms" => Some(1.0),
        "bohr" => Some(BOHR_TO_ANGSTROM),
        _ => None,
    }
}

fn parse_coordinate_component(token: &str, file: &str, context: &str) -> Result<f64> {
    token.parse::<f64>().map_err(|error| TbError::FileParse {
        file: file.to_string(),
        message: format!("Failed to parse {context}: {error}"),
    })
}

/// Trait for loading a tight-binding model from Wannier90 output files.
///
/// The implementing type controls which data is loaded:
/// - `DIM` must be `3` (Wannier90 always works in 3D).
/// - `R: RMatrixData` determines whether position matrix elements are loaded:
///   `HasRMatrix` requires `_r.dat` (generated by `write_rmn=true` in Wannier90),
///   while `NoRMatrix` skips it.
pub trait Wannier90 {
    fn from_hr(path: &str, file_name: &str, zero_energy: f64) -> Result<Self>
    where
        Self: Sized;
}

impl<const SPIN: bool, const DIM: usize, R: RMatrixData> Wannier90 for Model<SPIN, DIM, R> {
    #[allow(non_snake_case)]
    fn from_hr(path: &str, file_name: &str, zero_energy: f64) -> Result<Self> {
        // This function reads tight-binding files from Wannier90.
        //
        // The 'path' parameter specifies the file location, which can be an absolute path (starting with "/")
        // or a relative path relative to the directory when cargo run is executed.
        // The 'file_name' is the seedname in Wannier90. The function reads files:
        // seedname.win, seedname_centres.xyz, seedname_hr.dat, and optionally seedname_r.dat.
        //
        // For seedname_centres.xyz, set write_xyz=true in Wannier90; for seedname_hr.dat, set write_hr=true.
        //
        // Set write_rmn=true in Wannier90 to generate seedname_r.dat. This file is required when the
        // Model uses `HasRMatrix` (position matrix elements) for accurate velocity operators.
        // When `R = NoRMatrix`, the _r.dat file is skipped.
        //
        // DIM is always 3 for Wannier90 (the lattice matrix and all vectors are 3D).
        //
        // Additionally, for newer versions of Wannier90, to preserve good symmetry, it is recommended
        // to also provide wannier90_wsvec.dat for better symmetric results.

        use std::fs::File;
        use std::io::BufRead;
        use std::io::BufReader;
        use std::path::Path;

        let mut file_path = path.to_string();
        file_path.push_str(file_name);
        let mut hr_path = file_path.clone();
        hr_path.push_str("_hr.dat");

        let path = Path::new(&hr_path);
        let hr = File::open(path).map_err(|e| TbError::FileCreation {
            path: hr_path.clone(),
            message: format!("Unable to open HR file: {}", e),
        })?;
        let reader = BufReader::new(hr);
        let mut reads: Vec<String> = Vec::new();

        // 读取文件行
        for line in reader.lines() {
            let line = line.map_err(|e| TbError::Io(e))?;
            reads.push(line.clone());
        }

        // 获取轨道数和R点数
        let nsta = reads[1]
            .trim()
            .parse::<usize>()
            .map_err(|e| TbError::FileParse {
                file: hr_path.clone(),
                message: format!("Failed to parse nsta: {}", e),
            })?;
        let n_R = reads[2]
            .trim()
            .parse::<usize>()
            .map_err(|e| TbError::FileParse {
                file: hr_path.clone(),
                message: format!("Failed to parse n_R: {}", e),
            })?;
        let mut weights: Vec<usize> = Vec::new();
        let mut n_line: usize = 0;

        // 解析文件数据以获取权重
        for i in 3..reads.len() {
            if reads[i].contains(".") {
                n_line = i;
                break;
            }
            let string = reads[i].trim().split_whitespace();
            let string: Vec<_> = string
                .map(|x| {
                    x.parse::<usize>().map_err(|e| TbError::FileParse {
                        file: hr_path.clone(),
                        message: format!("Failed to parse weight: {}", e),
                    })
                })
                .collect::<Result<Vec<_>>>()?;
            weights.extend(string.clone());
        }

        // 初始化哈密顿量矩阵
        let mut hamR = Array2::<isize>::zeros((1, 3));
        let mut ham = Array3::<Complex<f64>>::zeros((1, nsta, nsta));

        // 遍历每个R点并填充哈密顿量
        for i in 0..n_R {
            let mut string = reads[i * nsta * nsta + n_line].trim().split_whitespace();
            let a = string
                .next()
                .ok_or_else(|| TbError::FileParse {
                    file: hr_path.clone(),
                    message: "Missing R vector component".to_string(),
                })?
                .parse::<isize>()
                .map_err(|e| TbError::FileParse {
                    file: hr_path.clone(),
                    message: format!("Failed to parse R vector: {}", e),
                })?;
            let b = string
                .next()
                .ok_or_else(|| TbError::FileParse {
                    file: hr_path.clone(),
                    message: "Missing R vector component".to_string(),
                })?
                .parse::<isize>()
                .map_err(|e| TbError::FileParse {
                    file: hr_path.clone(),
                    message: format!("Failed to parse R vector: {}", e),
                })?;
            let c = string
                .next()
                .ok_or_else(|| TbError::FileParse {
                    file: hr_path.clone(),
                    message: "Missing R vector component".to_string(),
                })?
                .parse::<isize>()
                .map_err(|e| TbError::FileParse {
                    file: hr_path.clone(),
                    message: format!("Failed to parse R vector: {}", e),
                })?;

            if a == 0 && b == 0 && c == 0 {
                for ind_i in 0..nsta {
                    for ind_j in 0..nsta {
                        let mut string = reads[i * nsta * nsta + ind_i * nsta + ind_j + n_line]
                            .trim()
                            .split_whitespace();
                        let re = string
                            .nth(5)
                            .ok_or_else(|| TbError::FileParse {
                                file: hr_path.clone(),
                                message: "Missing Hamiltonian real part".to_string(),
                            })?
                            .parse::<f64>()
                            .map_err(|e| TbError::FileParse {
                                file: hr_path.clone(),
                                message: format!("Failed to parse Hamiltonian real part: {}", e),
                            })?;
                        let im = string
                            .next()
                            .ok_or_else(|| TbError::FileParse {
                                file: hr_path.clone(),
                                message: "Missing Hamiltonian imaginary part".to_string(),
                            })?
                            .parse::<f64>()
                            .map_err(|e| TbError::FileParse {
                                file: hr_path.clone(),
                                message: format!(
                                    "Failed to parse Hamiltonian imaginary part: {}",
                                    e
                                ),
                            })?;
                        ham[[0, ind_j, ind_i]] = Complex::new(re, im) / (weights[i] as f64);
                    }
                }
            } else {
                let mut matrix = Array3::<Complex<f64>>::zeros((1, nsta, nsta));
                for ind_i in 0..nsta {
                    for ind_j in 0..nsta {
                        let mut string = reads[i * nsta * nsta + ind_i * nsta + ind_j + n_line]
                            .trim()
                            .split_whitespace();
                        let re = string
                            .nth(5)
                            .ok_or_else(|| TbError::FileParse {
                                file: hr_path.clone(),
                                message: "Missing Hamiltonian real part".to_string(),
                            })?
                            .parse::<f64>()
                            .map_err(|e| TbError::FileParse {
                                file: hr_path.clone(),
                                message: format!("Failed to parse Hamiltonian real part: {}", e),
                            })?;
                        let im = string
                            .next()
                            .ok_or_else(|| TbError::FileParse {
                                file: hr_path.clone(),
                                message: "Missing Hamiltonian imaginary part".to_string(),
                            })?
                            .parse::<f64>()
                            .map_err(|e| TbError::FileParse {
                                file: hr_path.clone(),
                                message: format!(
                                    "Failed to parse Hamiltonian imaginary part: {}",
                                    e
                                ),
                            })?;
                        matrix[[0, ind_j, ind_i]] = Complex::new(re, im) / (weights[i] as f64);
                        // wannier90 里面是按照纵向排列的矩阵
                    }
                }
                ham.append(Axis(0), matrix.view())
                    .map_err(|e| TbError::Linalg(ndarray_linalg::error::LinalgError::Shape(e)))?;
                hamR.append(Axis(0), arr2(&[[a, b, c]]).view())
                    .map_err(|e| TbError::Linalg(ndarray_linalg::error::LinalgError::Shape(e)))?;
            }
        }

        // 调整哈密顿量以匹配能量零点
        for i in 0..nsta {
            ham[[0, i, i]] -= Complex::new(zero_energy, 0.0);
        }
        //开始读取 .win 文件
        let _reads: Vec<String> = Vec::new();
        let mut win_path = file_path.clone();
        win_path.push_str(".win"); //文件的位置
        let path = Path::new(&win_path); //转化为路径格式
        let hr = File::open(path).map_err(|e| TbError::FileCreation {
            path: win_path.clone(),
            message: format!("Unable to open win file: {}", e),
        })?;
        let reader = BufReader::new(hr);
        let mut reads: Vec<String> = Vec::new();
        for line in reader.lines() {
            let line = line.map_err(|e| TbError::Io(e))?;
            reads.push(line.clone());
        }
        let mut read_iter = reads.iter();
        let mut lat = Array2::<f64>::zeros((3, 3)); //晶格轨道坐标初始化
        let mut spin: bool = false; //体系自旋初始化
        let _natom: usize = 0; //原子位置初始化
        let mut atom = Vec::new(); //原子位置坐标初始化
        let mut orb_proj = Vec::new();
        let mut proj_name = Vec::new();
        let mut proj_list: Vec<usize> = Vec::new();
        let _atom_list: Vec<usize> = Vec::new();
        let mut atom_name: Vec<&str> = Vec::new();
        let mut atom_pos = Array2::<f64>::zeros((0, 3));
        let mut atom_proj = Vec::new();
        loop {
            let a = read_iter.next();
            if a == None {
                break;
            } else {
                let a = a.ok_or_else(|| TbError::FileParse {
                    file: win_path.clone(),
                    message: "Unexpected end of file".to_string(),
                })?;
                let keyword = win_data_line(a).to_ascii_lowercase();
                if keyword.contains("begin unit_cell_cart") {
                    let mut unit_scale = 1.0_f64;
                    let mut rows = Vec::<[f64; 3]>::with_capacity(3);
                    while rows.len() < 3 {
                        let line = read_iter.next().ok_or_else(|| TbError::FileParse {
                            file: win_path.clone(),
                            message: "Missing lattice vector line".to_string(),
                        })?;
                        let line = win_data_line(line);
                        if line.is_empty() {
                            continue;
                        }
                        let tokens = line.split_whitespace().collect::<Vec<_>>();
                        if rows.is_empty() && tokens.len() == 1 {
                            unit_scale =
                                length_unit_scale(tokens[0]).ok_or_else(|| TbError::FileParse {
                                    file: win_path.clone(),
                                    message: format!("Unknown unit_cell_cart unit '{}'", tokens[0]),
                                })?;
                            continue;
                        }
                        if tokens.len() != 3 {
                            return Err(TbError::FileParse {
                                file: win_path.clone(),
                                message: format!(
                                    "A unit_cell_cart row must contain 3 numbers, found {} in '{line}'",
                                    tokens.len()
                                ),
                            });
                        }
                        rows.push([
                            parse_coordinate_component(tokens[0], &win_path, "lattice vector")?
                                * unit_scale,
                            parse_coordinate_component(tokens[1], &win_path, "lattice vector")?
                                * unit_scale,
                            parse_coordinate_component(tokens[2], &win_path, "lattice vector")?
                                * unit_scale,
                        ]);
                    }
                    for row in 0..3 {
                        for column in 0..3 {
                            lat[[row, column]] = rows[row][column];
                        }
                    }
                } else if keyword.contains("spinors")
                    && (keyword.contains('t') || keyword.contains("true"))
                {
                    spin = true;
                } else if keyword.contains("begin projections") {
                    loop {
                        let string = read_iter.next().ok_or_else(|| TbError::FileParse {
                            file: win_path.clone(),
                            message: "Unexpected end of file".to_string(),
                        })?;
                        let string = win_data_line(string);
                        if string.is_empty() {
                            continue;
                        }
                        if string.to_ascii_lowercase().contains("end projections") {
                            break;
                        } else {
                            let prj: Vec<&str> = string
                                .split(|c| c == ',' || c == ';' || c == ':')
                                .map(|x| x.trim())
                                .collect();
                            if prj.len() < 2 || prj[0].is_empty() {
                                return Err(TbError::FileParse {
                                    file: win_path.clone(),
                                    message: format!(
                                        "Malformed projection line '{string}': expected species:orbital"
                                    ),
                                });
                            }
                            let mut atom_orb_number: usize = 0;
                            let mut proj_orb = Vec::new();
                            for item in prj[1..].iter() {
                                let (aa, use_proj_orb): (usize, Vec<_>) = match (*item).trim() {
                                    "s" => (1, vec![OrbProj::s]),
                                    "p" => (3, vec![OrbProj::pz, OrbProj::px, OrbProj::py]),
                                    "d" => (
                                        5,
                                        vec![
                                            OrbProj::dz2,
                                            OrbProj::dxz,
                                            OrbProj::dyz,
                                            OrbProj::dx2y2,
                                            OrbProj::dxy,
                                        ],
                                    ),
                                    "f" => (
                                        7,
                                        vec![
                                            OrbProj::fz3,
                                            OrbProj::fxz2,
                                            OrbProj::fyz2,
                                            OrbProj::fzx2y2,
                                            OrbProj::fxyz,
                                            OrbProj::fxx23y2,
                                            OrbProj::fy3x2y2,
                                        ],
                                    ),
                                    "sp3" => (
                                        4,
                                        vec![
                                            OrbProj::sp3_1,
                                            OrbProj::sp3_2,
                                            OrbProj::sp3_3,
                                            OrbProj::sp3_4,
                                        ],
                                    ),
                                    "sp2" => {
                                        (3, vec![OrbProj::sp2_1, OrbProj::sp2_2, OrbProj::sp2_3])
                                    }
                                    "sp" => (2, vec![OrbProj::sp_1, OrbProj::sp_2]),
                                    "sp3d" => (
                                        5,
                                        vec![
                                            OrbProj::sp3d_1,
                                            OrbProj::sp3d_2,
                                            OrbProj::sp3d_3,
                                            OrbProj::sp3d_4,
                                            OrbProj::sp3d_5,
                                        ],
                                    ),
                                    "sp3d2" => (
                                        6,
                                        vec![
                                            OrbProj::sp3d2_1,
                                            OrbProj::sp3d2_2,
                                            OrbProj::sp3d2_3,
                                            OrbProj::sp3d2_4,
                                            OrbProj::sp3d2_5,
                                            OrbProj::sp3d2_6,
                                        ],
                                    ),
                                    "px" => (1, vec![OrbProj::px]),
                                    "py" => (1, vec![OrbProj::py]),
                                    "pz" => (1, vec![OrbProj::pz]),
                                    "dxy" => (1, vec![OrbProj::dxy]),
                                    "dxz" => (1, vec![OrbProj::dxz]),
                                    "dyz" => (1, vec![OrbProj::dyz]),
                                    "dz2" => (1, vec![OrbProj::dz2]),
                                    "dx2-y2" => (1, vec![OrbProj::dx2y2]),
                                    &_ => {
                                        return Err(TbError::InvalidOrbitalProjection(format!(
                                            "Unrecognized projection '{}' in seedname.win",
                                            item
                                        )));
                                    }
                                };
                                atom_orb_number += aa;
                                proj_orb.extend(use_proj_orb);
                            }
                            proj_list.push(atom_orb_number);
                            atom_proj.push(proj_orb);
                            let proj_type =
                                prj[0].parse::<AtomType>().map_err(|_| TbError::FileParse {
                                    file: win_path.clone(),
                                    message: format!(
                                        "Unknown atomic species '{}' in begin projections",
                                        prj[0]
                                    ),
                                })?;
                            proj_name.push(proj_type);
                        }
                    }
                } else if keyword.contains("begin atoms_cart") {
                    let mut cartesian_unit = 1.0_f64;
                    let mut first_data_line = true;
                    loop {
                        let string = read_iter.next().ok_or_else(|| TbError::FileParse {
                            file: win_path.clone(),
                            message: "Unexpected end of file".to_string(),
                        })?;
                        let string = win_data_line(string);
                        if string.is_empty() {
                            continue;
                        }
                        if string.to_ascii_lowercase().contains("end atoms_cart") {
                            break;
                        }
                        let fields = string.split_whitespace().collect::<Vec<_>>();
                        if first_data_line && fields.len() == 1 {
                            cartesian_unit =
                                length_unit_scale(fields[0]).ok_or_else(|| TbError::FileParse {
                                    file: win_path.clone(),
                                    message: format!("Unknown atoms_cart unit '{}'", fields[0]),
                                })?;
                            first_data_line = false;
                            continue;
                        }
                        first_data_line = false;
                        if fields.len() != 4 {
                            return Err(TbError::FileParse {
                                file: win_path.clone(),
                                message: format!(
                                    "An atoms_cart row must contain a species and 3 coordinates, found {} fields in '{string}'",
                                    fields.len()
                                ),
                            });
                        }
                        atom_name.push(fields[0]);
                        let position = array![
                            parse_coordinate_component(
                                fields[1],
                                &win_path,
                                "Cartesian atom position",
                            )? * cartesian_unit,
                            parse_coordinate_component(
                                fields[2],
                                &win_path,
                                "Cartesian atom position",
                            )? * cartesian_unit,
                            parse_coordinate_component(
                                fields[3],
                                &win_path,
                                "Cartesian atom position",
                            )? * cartesian_unit,
                        ];
                        // Only used when no _centres.xyz file is available.
                        atom_pos.push_row(position.view())?;
                    }
                } else if keyword.contains("begin atoms_frac") {
                    // Fractional positions; convert to Cartesian at parse
                    // time so the fallback path stays uniform.
                    loop {
                        let string = read_iter.next().ok_or_else(|| TbError::FileParse {
                            file: win_path.clone(),
                            message: "Unexpected end of file".to_string(),
                        })?;
                        let string = win_data_line(string);
                        if string.is_empty() {
                            continue;
                        }
                        if string.to_ascii_lowercase().contains("end atoms_frac") {
                            break;
                        }
                        let fields = string.split_whitespace().collect::<Vec<_>>();
                        if fields.len() != 4 {
                            return Err(TbError::FileParse {
                                file: win_path.clone(),
                                message: format!(
                                    "An atoms_frac row must contain a species and 3 coordinates, found {} fields in '{string}'",
                                    fields.len()
                                ),
                            });
                        }
                        atom_name.push(fields[0]);
                        let fractional = array![
                            parse_coordinate_component(
                                fields[1],
                                &win_path,
                                "fractional atom position",
                            )?,
                            parse_coordinate_component(
                                fields[2],
                                &win_path,
                                "fractional atom position",
                            )?,
                            parse_coordinate_component(
                                fields[3],
                                &win_path,
                                "fractional atom position",
                            )?,
                        ];
                        let position = fractional.dot(&lat);
                        atom_pos.push_row(position.view())?;
                    }
                }
            }
        }
        // 验证文件中的自旋设置与 SPIN 常量泛型是否一致
        if spin != SPIN {
            return Err(TbError::Other(format!(
                "Spin mismatch: Wannier90 .win file has spin={} but Model was constructed with SPIN={}",
                spin, SPIN
            )));
        }
        //开始读取 seedname_centres.xyz 文件
        let _reads: Vec<String> = Vec::new();
        let mut xyz_path = file_path.clone();
        xyz_path.push_str("_centres.xyz");
        let path = Path::new(&xyz_path);
        let hr = File::open(path);
        let orb = if let Ok(hr) = hr {
            let reader = BufReader::new(hr);
            let mut reads: Vec<String> = Vec::new();
            for line in reader.lines() {
                let line = line.map_err(|e| TbError::FileParse {
                    file: xyz_path.clone(),
                    message: format!("Failed to read line: {}", e),
                })?;
                reads.push(line.clone());
            }
            if reads.len() < 2 {
                return Err(TbError::FileParse {
                    file: xyz_path.clone(),
                    message: "_centres.xyz must contain a count and comment line".to_string(),
                });
            }
            let declared_entries =
                reads[0]
                    .trim()
                    .parse::<usize>()
                    .map_err(|error| TbError::FileParse {
                        file: xyz_path.clone(),
                        message: format!("Invalid _centres.xyz entry count: {error}"),
                    })?;
            let available_entries = reads.len() - 2;
            if declared_entries < nsta || available_entries < declared_entries {
                return Err(TbError::FileParse {
                    file: xyz_path.clone(),
                    message: format!(
                        "Truncated _centres.xyz: header declares {declared_entries} entries, \
                         HR requires {nsta} Wannier-centre entries, but only {} data lines are present",
                        available_entries
                    ),
                });
            }
            let norb = if spin { nsta / 2 } else { nsta };
            let mut orb = Array2::<f64>::zeros((norb, 3));
            for i in 0..norb {
                let fields = reads[i + 2].split_whitespace().collect::<Vec<_>>();
                if fields.len() < 4 {
                    return Err(TbError::FileParse {
                        file: xyz_path.clone(),
                        message: format!(
                            "Malformed Wannier-centre row {}: expected a label and 3 coordinates",
                            i + 1
                        ),
                    });
                }
                for axis in 0..3 {
                    orb[[i, axis]] = parse_coordinate_component(
                        fields[axis + 1],
                        &xyz_path,
                        "Wannier-centre position",
                    )?;
                }
            }
            orb = orb.dot(&lat.inv().map_err(TbError::Linalg)?);
            let atom_count = declared_entries - nsta;
            let mut new_atom_pos = Array2::<f64>::zeros((atom_count, 3));
            let mut new_atom_name = Vec::with_capacity(atom_count);
            for i in 0..atom_count {
                let fields = reads[i + 2 + nsta].split_whitespace().collect::<Vec<_>>();
                if fields.len() < 4 {
                    return Err(TbError::FileParse {
                        file: xyz_path.clone(),
                        message: format!(
                            "Malformed atom row {} in _centres.xyz: expected a species and 3 coordinates",
                            i + 1
                        ),
                    });
                }
                for axis in 0..3 {
                    new_atom_pos[[i, axis]] =
                        parse_coordinate_component(fields[axis + 1], &xyz_path, "atom position")?;
                }
                let name = fields[0]
                    .parse::<AtomType>()
                    .map_err(|_| TbError::FileParse {
                        file: xyz_path.clone(),
                        message: format!(
                            "Unknown atomic species '{}' in _centres.xyz; \
                         species must be listed in the win file's begin projections block",
                            fields[0]
                        ),
                    })?;
                new_atom_name.push(name);
            }
            //接下来如果wannier90.win 和 .xyz 文件的原子顺序不一致, 那么我们以xyz的原子顺序为准, 调整 atom_list

            let mut dropped: Vec<AtomType> = Vec::new();
            for (i, name) in new_atom_name.iter().enumerate() {
                // Multiple projection lines of the same species belong to the
                // SAME atom; merge them into one Atom.
                let mut atom_orbitals = Vec::new();
                for (j, j_name) in proj_name.iter().enumerate() {
                    if j_name == name {
                        let first = orb_proj.len();
                        atom_orbitals.extend((first..first + proj_list[j]).map(OrbitalId::new));
                        orb_proj.extend(atom_proj[j].clone());
                    }
                }
                if atom_orbitals.is_empty() {
                    // 该物种在 projections 块中没有被拟合任何轨道, 默认直接丢弃该物种的所有原子,
                    // 而不是报错 (例如只提供了 Cs 的坐标但没有在 Cs 上拟合 Wannier 轨道)。
                    // 每个物种只警告一次, 避免同一物种的多个原子重复刷屏。
                    if !dropped.contains(name) {
                        eprintln!(
                            "warning: dropping '{}' atoms from _centres.xyz (no orbitals in the projections block)",
                            name.to_str()
                        );
                        dropped.push(*name);
                    }
                    continue;
                }
                let use_pos = new_atom_pos
                    .row(i)
                    .dot(&lat.inv().map_err(TbError::Linalg)?);
                atom.push(Atom::with_orbitals(use_pos, *name, atom_orbitals));
            }
            // 所有轨道都必须能被 xyz 物种与 projections 物种的匹配覆盖,
            // 否则 orb_proj 条目数少于 norb, validate() 会报 orbital_projection_count。
            if orb_proj.len() != norb {
                let xyz_species: Vec<&str> = new_atom_name.iter().map(|n| n.to_str()).collect();
                let proj_species: Vec<&str> = proj_name.iter().map(|n| n.to_str()).collect();
                let detail = if orb_proj.len() < norb {
                    format!(
                        "{norb} orbitals declared in the projections block, but only {} could be assigned to atoms",
                        orb_proj.len()
                    )
                } else {
                    format!(
                        "{norb} orbitals declared in the projections block, but {} were assigned to atoms",
                        orb_proj.len()
                    )
                };
                return Err(TbError::FileParse {
                    file: xyz_path.clone(),
                    message: format!(
                        "species mismatch between _centres.xyz and the win projections block: \
                         {detail}. _centres.xyz species: {xyz_species:?}; \
                         projection species: {proj_species:?}"
                    ),
                });
            }
            orb
        } else {
            let mut orb = Array2::<f64>::zeros((0, 3));
            let atom_pos = atom_pos.dot(&lat.inv().map_err(TbError::Linalg)?);
            let mut dropped: Vec<AtomType> = Vec::new();
            for (i, name) in atom_name.iter().enumerate() {
                let name = name.parse::<AtomType>().map_err(|_| TbError::FileParse {
                    file: win_path.clone(),
                    message: format!("Unknown atomic species '{name}' in begin atoms_frac block"),
                })?;
                // Wannier90 permits multiple projection lines per species
                // (e.g. separate spin-up/down blocks): all lines of this
                // species belong to the SAME atom, so merge them into one
                // Atom instead of creating one Atom per line.
                let mut atom_orbitals = Vec::new();
                for (j, j_name) in proj_name.iter().enumerate() {
                    if name == *j_name {
                        let first = orb_proj.len();
                        atom_orbitals.extend((first..first + proj_list[j]).map(OrbitalId::new));
                        orb_proj.extend(atom_proj[j].clone());
                        for _ in 0..proj_list[j] {
                            orb.push_row(atom_pos.row(i).view())?;
                        }
                    }
                }
                if atom_orbitals.is_empty() {
                    // 该物种在 projections 块中没有被拟合任何轨道, 默认直接丢弃该物种的所有原子。
                    // 每个物种只警告一次, 避免同一物种的多个原子重复刷屏。
                    if !dropped.contains(&name) {
                        eprintln!(
                            "warning: dropping '{name}' atoms from begin atoms_frac/atoms_cart (no orbitals in the projections block)"
                        );
                        dropped.push(name);
                    }
                    continue;
                }
                atom.push(Atom::with_orbitals(
                    atom_pos.row(i).to_owned(),
                    name,
                    atom_orbitals,
                ));
            }
            // 与 xyz 分支相同的物种覆盖检查: atoms_frac 的物种必须全部出现在
            // projections 块中, 否则 orb_proj 条目数不足, validate() 会失败。
            // The HR file's nsta is authoritative: every atom of a species
            // gets its own copy of each projection line, so the constructed
            // count is (projection lines x atoms), not the once-only win
            // declaration count.
            let expected_norb = if spin { nsta / 2 } else { nsta };
            if orb_proj.len() != expected_norb {
                let proj_species: Vec<&str> = proj_name.iter().map(|n| n.to_str()).collect();
                let detail = if orb_proj.len() < expected_norb {
                    format!(
                        "{expected_norb} orbitals declared in the HR file, but only {} could be assigned to atoms",
                        orb_proj.len()
                    )
                } else {
                    format!(
                        "{expected_norb} orbitals declared in the HR file, but {} were assigned to atoms",
                        orb_proj.len()
                    )
                };
                return Err(TbError::FileParse {
                    file: win_path.clone(),
                    message: format!(
                        "species mismatch between begin atoms_frac and the projections block: \
                         {detail}. atom species: {atom_name:?}; \
                         projection species: {proj_species:?}"
                    ),
                });
            }
            orb
        };
        //开始尝试读取 _r.dat 文件
        let mut have_r = false;
        let mut rmatrix = if R::HAS_RMATRIX {
            let mut r_path = file_path.clone();
            r_path.push_str("_r.dat");
            let path = Path::new(&r_path);
            let hr = File::open(path);
            if let Ok(hr) = hr {
                have_r = true;
                let reader = BufReader::new(hr);
                let mut reads: Vec<String> = Vec::new();
                for line in reader.lines() {
                    let line = line.map_err(|e| TbError::FileParse {
                        file: xyz_path.clone(),
                        message: format!("Failed to read line: {}", e),
                    })?;
                    reads.push(line.clone());
                }
                let n_R = reads[2]
                    .trim()
                    .parse::<usize>()
                    .map_err(|e| TbError::FileParse {
                        file: r_path.clone(),
                        message: format!("Failed to parse n_R: {}", e),
                    })?;
                let mut rmatrix = Array4::<Complex<f64>>::zeros((hamR.nrows(), 3, nsta, nsta));
                for i in 0..n_R {
                    let mut string = reads[i * nsta * nsta + 3].trim().split_whitespace();
                    let a = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: r_path.clone(),
                            message: "Missing R vector component".to_string(),
                        })?
                        .parse::<isize>()
                        .map_err(|e| TbError::FileParse {
                            file: r_path.clone(),
                            message: format!("Failed to parse R vector: {}", e),
                        })?;
                    let b = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: r_path.clone(),
                            message: "Missing R vector component".to_string(),
                        })?
                        .parse::<isize>()
                        .map_err(|e| TbError::FileParse {
                            file: r_path.clone(),
                            message: format!("Failed to parse R vector: {}", e),
                        })?;
                    let c = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: r_path.clone(),
                            message: "Missing R vector component".to_string(),
                        })?
                        .parse::<isize>()
                        .map_err(|e| TbError::FileParse {
                            file: r_path.clone(),
                            message: format!("Failed to parse R vector: {}", e),
                        })?;
                    let R0 = array![a, b, c];
                    let index = find_R(&hamR, &R0).ok_or_else(|| TbError::FileParse {
                        file: r_path.clone(),
                        message: format!("R vector {:?} not found in Hamiltonian", R0),
                    })?;
                    for ind_i in 0..nsta {
                        for ind_j in 0..nsta {
                            let string = &reads[i * nsta * nsta + ind_i * nsta + ind_j + 3];
                            let mut string = string.trim().split_whitespace();
                            string.nth(4);
                            for r in 0..3 {
                                let re = string
                                    .next()
                                    .ok_or_else(|| TbError::FileParse {
                                        file: r_path.clone(),
                                        message: "Missing R matrix real part".to_string(),
                                    })?
                                    .parse::<f64>()
                                    .map_err(|e| TbError::FileParse {
                                        file: r_path.clone(),
                                        message: format!(
                                            "Failed to parse R matrix real part: {}",
                                            e
                                        ),
                                    })?;
                                let im = string
                                    .next()
                                    .ok_or_else(|| TbError::FileParse {
                                        file: r_path.clone(),
                                        message: "Missing R matrix imaginary part".to_string(),
                                    })?
                                    .parse::<f64>()
                                    .map_err(|e| TbError::FileParse {
                                        file: r_path.clone(),
                                        message: format!(
                                            "Failed to parse R matrix imaginary part: {}",
                                            e
                                        ),
                                    })?;
                                rmatrix[[index, r, ind_j, ind_i]] =
                                    Complex::new(re, im) / (weights[i] as f64);
                            }
                        }
                    }
                }
                rmatrix
            } else {
                return Err(TbError::FileCreation {
                    path: r_path.clone(),
                    message: "R::HAS_RMATRIX=true but _r.dat file not found".to_string(),
                });
            }
        } else {
            Array4::<Complex<f64>>::zeros((1, 3, 1, 1))
        };

        //最后判断有没有wannier90_wsvec.dat-----------------------------------
        let mut ws_path = file_path.clone();
        ws_path.push_str("_wsvec.dat");
        let path = Path::new(&ws_path); //转化为路径格式
        let ws = File::open(path);
        if let Ok(ws) = ws {
            let reader = BufReader::new(ws);
            let mut reads: Vec<String> = Vec::new();
            for line in reader.lines() {
                let line = line.map_err(|e| TbError::FileParse {
                    file: xyz_path.clone(),
                    message: format!("Failed to read line: {}", e),
                })?;
                reads.push(line.clone());
            }
            //开始针对ham, hamR 以及 rmatrix 进行修改
            //我们先考虑有rmatrix的情况
            if have_r {
                let mut i = 0;
                let mut new_hamR = Array2::zeros((1, 3));
                let mut new_ham = Array3::zeros((1, nsta, nsta));
                let mut new_rmatrix = Array4::zeros((1, 3, nsta, nsta));
                while i < reads.len() - 1 {
                    i += 1;
                    let line = &reads[i];
                    let mut string = line.trim().split_whitespace();
                    let a = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing R vector component".to_string(),
                        })?
                        .parse::<isize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse R vector: {}", e),
                        })?;
                    let b = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing R vector component".to_string(),
                        })?
                        .parse::<isize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse R vector: {}", e),
                        })?;
                    let c = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing R vector component".to_string(),
                        })?
                        .parse::<isize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse R vector: {}", e),
                        })?;
                    let int_i = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing orbital index".to_string(),
                        })?
                        .parse::<usize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse orbital index: {}", e),
                        })?
                        - 1;
                    let int_j = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing orbital index".to_string(),
                        })?
                        .parse::<usize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse orbital index: {}", e),
                        })?
                        - 1;
                    //接下来判断是否在我们的hamR 中
                    i += 1;
                    let weight = reads[i]
                        .trim()
                        .split_whitespace()
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing weight value".to_string(),
                        })?
                        .parse::<usize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse weight: {}", e),
                        })?;
                    let R = array![a, b, c];
                    let index = find_R(&hamR, &R).ok_or_else(|| TbError::FileParse {
                        file: ws_path.clone(),
                        message: format!("R vector {:?} not found in Hamiltonian", R),
                    })?;
                    let hop = ham[[index, int_i, int_j]] / (weight as f64);
                    let hop_x = rmatrix[[index, 0, int_i, int_j]] / (weight as f64);
                    let hop_y = rmatrix[[index, 1, int_i, int_j]] / (weight as f64);
                    let hop_z = rmatrix[[index, 2, int_i, int_j]] / (weight as f64);

                    for _i0 in 0..weight {
                        i += 1;
                        let line = &reads[i];
                        let mut string = line.trim().split_whitespace();
                        let a = string.next().unwrap().parse::<isize>().unwrap();
                        let b = string.next().unwrap().parse::<isize>().unwrap();
                        let c = string.next().unwrap().parse::<isize>().unwrap();
                        let new_R = array![R[[0]] + a, R[[1]] + b, R[[2]] + c];
                        if let Some(index0) = find_R(&new_hamR, &new_R) {
                            new_ham[[index0, int_i, int_j]] += hop;
                            new_rmatrix[[index0, 0, int_i, int_j]] += hop_x;
                            new_rmatrix[[index0, 1, int_i, int_j]] += hop_y;
                            new_rmatrix[[index0, 2, int_i, int_j]] += hop_z;
                        } else {
                            let mut use_ham = Array2::zeros((nsta, nsta));
                            let mut use_rmatrix = Array3::zeros((3, nsta, nsta));
                            use_ham[[int_i, int_j]] += hop;
                            use_rmatrix[[0, int_i, int_j]] += hop_x;
                            use_rmatrix[[1, int_i, int_j]] += hop_y;
                            use_rmatrix[[2, int_i, int_j]] += hop_z;
                            new_hamR.push_row(new_R.view())?;
                            new_ham.push(Axis(0), use_ham.view())?;
                            new_rmatrix.push(Axis(0), use_rmatrix.view())?;
                        }
                    }
                }
                hamR = new_hamR;
                ham = new_ham;
                rmatrix = new_rmatrix;
            } else {
                let mut i = 0;
                let mut new_hamR = Array2::zeros((1, 3));
                let mut new_ham = Array3::zeros((1, nsta, nsta));
                while i < reads.len() - 1 {
                    i += 1;
                    let line = &reads[i];
                    let mut string = line.trim().split_whitespace();
                    let a = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing R vector component".to_string(),
                        })?
                        .parse::<isize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse R vector: {}", e),
                        })?;
                    let b = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing R vector component".to_string(),
                        })?
                        .parse::<isize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse R vector: {}", e),
                        })?;
                    let c = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing R vector component".to_string(),
                        })?
                        .parse::<isize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse R vector: {}", e),
                        })?;
                    let int_i = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing orbital index".to_string(),
                        })?
                        .parse::<usize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse orbital index: {}", e),
                        })?
                        - 1;
                    let int_j = string
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing orbital index".to_string(),
                        })?
                        .parse::<usize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse orbital index: {}", e),
                        })?
                        - 1;
                    //接下来判断是否在我们的hamR 中
                    i += 1;
                    let weight = reads[i]
                        .trim()
                        .split_whitespace()
                        .next()
                        .ok_or_else(|| TbError::FileParse {
                            file: ws_path.clone(),
                            message: "Missing weight value".to_string(),
                        })?
                        .parse::<usize>()
                        .map_err(|e| TbError::FileParse {
                            file: ws_path.clone(),
                            message: format!("Failed to parse weight: {}", e),
                        })?;
                    let R = array![a, b, c];
                    let index = find_R(&hamR, &R).ok_or_else(|| TbError::FileParse {
                        file: ws_path.clone(),
                        message: format!("R vector {:?} not found in Hamiltonian", R),
                    })?;
                    let hop = ham[[index, int_i, int_j]] / (weight as f64);

                    for _i0 in 0..weight {
                        i += 1;
                        let line = &reads[i];
                        let mut string = line.trim().split_whitespace();
                        let a = string.next().unwrap().parse::<isize>().unwrap();
                        let b = string.next().unwrap().parse::<isize>().unwrap();
                        let c = string.next().unwrap().parse::<isize>().unwrap();
                        let new_R = array![R[[0]] + a, R[[1]] + b, R[[2]] + c];
                        if let Some(index0) = find_R(&new_hamR, &new_R) {
                            new_ham[[index0, int_i, int_j]] += hop;
                        } else {
                            let mut use_ham = Array2::zeros((nsta, nsta));
                            use_ham[[int_i, int_j]] = hop;
                            new_hamR.push_row(new_R.view())?;
                            new_ham.push(Axis(0), use_ham.view())?;
                        }
                    }
                }
                hamR = new_hamR;
                ham = new_ham;
            }
        }
        //最后一步, 将rmatrix 变成厄密的

        if have_r {
            for r in 0..hamR.nrows() - 1 {
                let R = hamR.row(r);
                let R_inv = -&R;
                if let Some(index) = find_R(&hamR, &R_inv) {
                    for i in 0..nsta {
                        for j in 0..nsta {
                            rmatrix[[r, 0, i, j]] =
                                (rmatrix[[r, 0, i, j]] + rmatrix[[index, 0, j, i]].conj()) / 2.0;
                            rmatrix[[r, 1, i, j]] =
                                (rmatrix[[r, 1, i, j]] + rmatrix[[index, 1, j, i]].conj()) / 2.0;
                            rmatrix[[r, 2, i, j]] =
                                (rmatrix[[r, 2, i, j]] + rmatrix[[index, 2, j, i]].conj()) / 2.0;
                            rmatrix[[index, 0, j, i]] = rmatrix[[r, 0, i, j]].conj();
                            rmatrix[[index, 1, j, i]] = rmatrix[[r, 1, i, j]].conj();
                            rmatrix[[index, 2, j, i]] = rmatrix[[r, 2, i, j]].conj();
                        }
                    }
                } else {
                    return Err(TbError::MissingHermitianConjugate { r: R.to_owned() });
                }
            }
        }

        // Validate that loaded data dimension matches the const generic DIM
        if DIM != 3 {
            return Err(TbError::InvalidDimension {
                dim: DIM,
                supported: vec![3],
            });
        }
        let model = Self {
            lat,
            orb,
            orb_projection: orb_proj,
            atoms: atom,
            ham,
            hamR,
            rmatrix: R::from_array(rmatrix),
        };
        model.validate()?;
        Ok(model)
    }
}

impl<const SPIN: bool, const DIM: usize> Model<SPIN, DIM, HasRMatrix> {
    /// Load a tight-binding model from Wannier90 files including position matrix elements.
    ///
    /// This is a convenience wrapper around [`Wannier90::from_hr`] that requires
    /// the Wannier90 `_r.dat` file to be present. Returns `Model<SPIN, DIM, HasRMatrix>`
    /// with position matrix elements (`rmatrix`) populated.
    ///
    /// # Errors
    ///
    /// Returns `TbError::FileCreation` if `_r.dat` is missing.
    pub fn from_hr_with_rmatrix(path: &str, file_name: &str, zero_energy: f64) -> Result<Self> {
        <Self as Wannier90>::from_hr(path, file_name, zero_energy)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use std::io::Write;

    /// Write a minimal valid Wannier90 dataset (one C atom, one s orbital)
    /// to `dir/seedname.*`, returning the directory name.
    fn write_minimal_dataset(dir: &str, atom_species: &str) {
        fs::create_dir_all(dir).unwrap();
        let win = format!(
            "begin unit_cell_cart\n1.0 0.0 0.0\n0.0 1.0 0.0\n0.0 0.0 1.0\nend unit_cell_cart\n\nbegin projections\nC:s\nend projections\n"
        );
        let hr = "generated\n1\n1\n1\n0 0 0 1 1 0.0 0.0\n";
        let xyz = format!("2\nWannier centres\nC 0.0 0.0 0.0\n{atom_species} 0.0 0.0 0.0\n");
        for (suffix, content) in [("_hr.dat", hr), ("_centres.xyz", xyz.as_str())] {
            let mut f = fs::File::create(format!("{dir}/seedname{suffix}")).unwrap();
            f.write_all(content.as_bytes()).unwrap();
        }
        let mut f = fs::File::create(format!("{dir}/seedname.win")).unwrap();
        f.write_all(win.as_bytes()).unwrap();
    }

    #[test]
    fn from_hr_loads_minimal_dataset() {
        let dir = "tests/tmp_w90_ok/";
        write_minimal_dataset(dir, "C");
        let model = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap();
        assert_eq!(model.norb(), 1);
        assert_eq!(model.natom(), 1);
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn from_hr_drops_xyz_atom_without_orbitals() {
        // Regression: an xyz atom whose species has no projection (e.g. a Cs
        // site that was never Wannierized) has no fitted orbitals, so it must
        // be dropped while atoms with orbitals are kept.
        let dir = "tests/tmp_w90_drop/";
        let _ = fs::remove_dir_all(dir);
        fs::create_dir_all(dir).unwrap();
        let win = "begin unit_cell_cart\n1.0 0.0 0.0\n0.0 1.0 0.0\n0.0 0.0 1.0\nend unit_cell_cart\n\nbegin projections\nC:s\nend projections\n";
        let hr = "generated\n1\n1\n1\n0 0 0 1 1 0.0 0.0\n";
        // One Wannier centre, one C atom with a fitted orbital, one Cs atom
        // with no orbitals: the Cs atom must be dropped.
        let xyz = "3\nWannier centres\nX 0.0 0.0 0.0\nC 0.1 0.0 0.0\nCs 0.2 0.0 0.0\n";
        fs::write(format!("{dir}seedname.win"), win).unwrap();
        fs::write(format!("{dir}seedname_hr.dat"), hr).unwrap();
        fs::write(format!("{dir}seedname_centres.xyz"), xyz).unwrap();

        let model = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap();
        assert_eq!(model.norb(), 1);
        assert_eq!(
            model.natom(),
            1,
            "Cs atom has no orbitals and must be dropped"
        );
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn from_hr_drops_multiple_atoms_of_same_unprojected_species() {
        // Two Cs atoms, both without orbitals: both must be dropped (the
        // warning is emitted once per species, but that is stderr-only and
        // not asserted here).
        let dir = "tests/tmp_w90_drop_multi/";
        let _ = fs::remove_dir_all(dir);
        fs::create_dir_all(dir).unwrap();
        let win = "begin unit_cell_cart\n1.0 0.0 0.0\n0.0 1.0 0.0\n0.0 0.0 1.0\nend unit_cell_cart\n\nbegin projections\nC:s\nend projections\n";
        let hr = "generated\n1\n1\n1\n0 0 0 1 1 0.0 0.0\n";
        let xyz =
            "4\nWannier centres\nX 0.0 0.0 0.0\nC 0.1 0.0 0.0\nCs 0.2 0.0 0.0\nCs 0.3 0.0 0.0\n";
        fs::write(format!("{dir}seedname.win"), win).unwrap();
        fs::write(format!("{dir}seedname_hr.dat"), hr).unwrap();
        fs::write(format!("{dir}seedname_centres.xyz"), xyz).unwrap();

        let model = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap();
        assert_eq!(model.norb(), 1);
        assert_eq!(model.natom(), 1, "both Cs atoms must be dropped");
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn from_hr_without_xyz_drops_atom_without_orbitals() {
        // No _centres.xyz: the atoms_cart fallback must also drop a species
        // with no projection instead of erroring.
        let dir = "tests/tmp_w90_no_xyz_drop/";
        let _ = fs::remove_dir_all(dir);
        fs::create_dir_all(dir).unwrap();
        let win = "begin unit_cell_cart\n1.0 0.0 0.0\n0.0 1.0 0.0\n0.0 0.0 1.0\nend unit_cell_cart\n\nbegin atoms_cart\nC 0.0 0.0 0.0\nCs 0.1 0.0 0.0\nend atoms_cart\n\nbegin projections\nC:s\nend projections\n";
        let hr = "generated\n1\n1\n1\n0 0 0 1 1 0.0 0.0\n";
        fs::write(format!("{dir}seedname.win"), win).unwrap();
        fs::write(format!("{dir}seedname_hr.dat"), hr).unwrap();

        let model = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap();
        assert_eq!(model.norb(), 1);
        assert_eq!(model.natom(), 1, "Cs atom must be dropped");
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn from_hr_rejects_projection_species_missing_from_xyz() {
        // Reverse direction: the projection block declares C:s, but no C atom
        // appears in _centres.xyz (only Fe, which is dropped). The C orbital
        // can never be assigned, so this must remain a hard error.
        let dir = "tests/tmp_w90_proj_missing/";
        write_minimal_dataset(dir, "Fe");
        let err = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap_err();
        assert!(matches!(err, TbError::FileParse { .. }));
        assert!(
            err.to_string().contains("species mismatch"),
            "unexpected error: {err}"
        );
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn from_hr_rejects_unknown_projection_species() {
        // Regression: `parse::<AtomType>()` must reject an unknown species in
        // the projections block with the same FileParse error as before.
        let dir = "tests/tmp_w90_bad_proj_species/";
        let _ = fs::remove_dir_all(dir);
        fs::create_dir_all(dir).unwrap();
        let win = "begin unit_cell_cart\n1.0 0.0 0.0\n0.0 1.0 0.0\n0.0 0.0 1.0\nend unit_cell_cart\n\nbegin projections\nXx:s\nend projections\n";
        let hr = "generated\n1\n1\n1\n0 0 0 1 1 0.0 0.0\n";
        let xyz = "1\nWannier centres\nC 0.0 0.0 0.0\n";
        fs::write(format!("{dir}seedname.win"), win).unwrap();
        fs::write(format!("{dir}seedname_hr.dat"), hr).unwrap();
        fs::write(format!("{dir}seedname_centres.xyz"), xyz).unwrap();

        let err = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap_err();
        assert!(matches!(err, TbError::FileParse { .. }));
        assert!(
            err.to_string()
                .contains("Unknown atomic species 'Xx' in begin projections"),
            "unexpected error: {err}"
        );
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn win_data_line_strips_hash_and_bang_comments() {
        assert_eq!(win_data_line("spinors = .true."), "spinors = .true.");
        assert_eq!(win_data_line("# spinors = .true."), "");
        assert_eq!(win_data_line("#spinors=true"), "");
        assert_eq!(
            win_data_line("spinors = .true. ! enable spin"),
            "spinors = .true."
        );
        assert_eq!(win_data_line("Fe 1.0 0.0 0.0 # position"), "Fe 1.0 0.0 0.0");
        assert_eq!(win_data_line("  \t"), "");
    }

    #[test]
    fn from_hr_without_centres_xyz_uses_atoms_cart_fallback() {
        // Regression: without _centres.xyz the outer norb stayed 0 and the
        // atoms_cart fallback underflowed (0 - 1) in debug builds.
        let dir = "tests/tmp_w90_no_xyz/";
        fs::create_dir_all(dir).unwrap();
        let win = "begin unit_cell_cart\n1.0 0.0 0.0\n0.0 1.0 0.0\n0.0 0.0 1.0\nend unit_cell_cart\n\nbegin atoms_cart\nFe 0.0 0.0 0.0\nend atoms_cart\n\nbegin projections\nFe:s\nend projections\n";
        let hr = "generated\n1\n1\n1\n0 0 0 1 1 0.0 0.0\n";
        let mut f = fs::File::create(format!("{dir}seedname.win")).unwrap();
        f.write_all(win.as_bytes()).unwrap();
        let mut f = fs::File::create(format!("{dir}seedname_hr.dat")).unwrap();
        f.write_all(hr.as_bytes()).unwrap();
        let model = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap();
        assert_eq!(model.norb(), 1);
        assert_eq!(model.natom(), 1);
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn from_hr_without_xyz_loads_multiple_atoms_of_same_species() {
        // Regression: the no-xyz fallback compared the constructed orbital
        // count against the once-only win projection declaration (4 for
        // Fe:s,p), but two Fe atoms each get their own copy (8 orbitals),
        // which the HR file confirms via nsta.
        let dir = "tests/tmp_w90_multi_atom/";
        fs::create_dir_all(dir).unwrap();
        let win = "begin unit_cell_cart\n1.0 0.0 0.0\n0.0 1.0 0.0\n0.0 0.0 1.0\nend unit_cell_cart\n\nbegin atoms_cart\nFe 0.0 0.0 0.0\nFe 1.0 0.0 0.0\nend atoms_cart\n\nbegin projections\nFe:s\nFe:p\nend projections\n";
        // 8 orbitals: 8x8 identity Hamiltonian at R=0.
        let mut hr = String::from("generated\n8\n1\n1\n0 0 0 1 1 0.0 0.0\n");
        for i in 1..=8 {
            for j in 1..=8 {
                if i != 1 || j != 1 {
                    let value = if i == j { "0.0 0.0" } else { "0.0 0.0" };
                    hr.push_str(&format!("0 0 0 {i} {j} {value}\n"));
                }
            }
        }
        let mut f = fs::File::create(format!("{dir}seedname.win")).unwrap();
        f.write_all(win.as_bytes()).unwrap();
        let mut f = fs::File::create(format!("{dir}seedname_hr.dat")).unwrap();
        f.write_all(hr.as_bytes()).unwrap();
        let model = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap();
        assert_eq!(model.norb(), 8);
        assert_eq!(model.natom(), 2);
        assert_eq!(model.atoms[0].norb(), 4);
        assert_eq!(model.atoms[1].norb(), 4);
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn from_hr_handles_coordinate_blocks_and_units_consistently() {
        // Regression: unit_cell_cart treated its optional unit as the first
        // lattice row, atoms_cart only accepted lower-case units, and
        // atoms_frac was not parsed at all. Each case below describes the same
        // atom at fractional x=0.5 in a two-unit cubic cell.
        for (lattice_unit, block, expected_lattice) in [
            (
                "BoHr",
                "ATOMS_CART\nBOHR\nFe 1.0 0.0 0.0 ! inline comment\nEND ATOMS_CART",
                2.0 * BOHR_TO_ANGSTROM,
            ),
            (
                "AnG",
                "atoms_cart\nAngstrom\nFe 1.0 0.0 0.0\nend atoms_cart",
                2.0,
            ),
            ("Ang", "atoms_frac\nFe 0.5 0.0 0.0\nend atoms_frac", 2.0),
        ] {
            let dir = "tests/tmp_w90_units/";
            let _ = fs::remove_dir_all(dir);
            fs::create_dir_all(dir).unwrap();
            let win = format!(
                "BEGIN UNIT_CELL_CART\n{lattice_unit}\n2.0 0.0 0.0\n0.0 2.0 0.0\n0.0 0.0 2.0\nEND UNIT_CELL_CART\n\nBEGIN {block}\n\nbegin projections\nFe:s\nend projections\n"
            );
            let hr = "generated\n1\n1\n1\n0 0 0 1 1 0.0 0.0\n";
            let mut f = fs::File::create(format!("{dir}seedname.win")).unwrap();
            f.write_all(win.as_bytes()).unwrap();
            let mut f = fs::File::create(format!("{dir}seedname_hr.dat")).unwrap();
            f.write_all(hr.as_bytes()).unwrap();
            let model = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap();
            assert_eq!(model.norb(), 1);
            assert_eq!(model.natom(), 1);
            assert!((model.lat[[0, 0]] - expected_lattice).abs() < 1e-12);
            assert!((model.atoms[0].position()[0] - 0.5).abs() < 1e-12);
            assert!((model.orb[[0, 0]] - 0.5).abs() < 1e-12);
        }
        let _ = fs::remove_dir_all("tests/tmp_w90_units");
    }

    #[test]
    fn from_hr_rejects_malformed_atom_rows_without_panicking() {
        let dir = "tests/tmp_w90_malformed_atom/";
        let _ = fs::remove_dir_all(dir);
        fs::create_dir_all(dir).unwrap();
        let win = "begin unit_cell_cart\n1 0 0\n0 1 0\n0 0 1\nend unit_cell_cart\n\nbegin atoms_cart\nFe 0.0 0.0\nend atoms_cart\n\nbegin projections\nFe:s\nend projections\n";
        let hr = "generated\n1\n1\n1\n0 0 0 1 1 0.0 0.0\n";
        fs::write(format!("{dir}seedname.win"), win).unwrap();
        fs::write(format!("{dir}seedname_hr.dat"), hr).unwrap();

        let result = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0);
        assert!(matches!(result, Err(TbError::FileParse { .. })));
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn from_hr_rejects_truncated_centres_xyz_without_panicking() {
        let dir = "tests/tmp_w90_truncated_xyz/";
        let _ = fs::remove_dir_all(dir);
        fs::create_dir_all(dir).unwrap();
        let win = "begin unit_cell_cart\n1 0 0\n0 1 0\n0 0 1\nend unit_cell_cart\n\nbegin projections\nFe:s,p\nend projections\n";
        let mut hr = String::from("generated\n4\n1\n1\n");
        for i in 1..=4 {
            for j in 1..=4 {
                hr.push_str(&format!("0 0 0 {i} {j} 0.0 0.0\n"));
            }
        }
        // The header declares four centres plus one atom, but only one centre
        // and one atom line are present.
        let xyz = "5\nWannier centres\nX 0.0 0.0 0.0\nFe 0.0 0.0 0.0\n";
        fs::write(format!("{dir}seedname.win"), win).unwrap();
        fs::write(format!("{dir}seedname_hr.dat"), hr).unwrap();
        fs::write(format!("{dir}seedname_centres.xyz"), xyz).unwrap();

        let error = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap_err();
        assert!(matches!(&error, TbError::FileParse { .. }));
        assert!(error.to_string().contains("Truncated"));
        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn from_hr_merges_multi_line_projections_into_one_atom() {
        // Regression: multiple projection lines of the same species (e.g.
        // Fe:s and Fe:p) created one Atom per line; Wannier90 supports
        // multiple lines per site and they must merge into a single Atom.
        let dir = "tests/tmp_w90_multiline/";
        fs::create_dir_all(dir).unwrap();
        let win = "begin unit_cell_cart\n1.0 0.0 0.0\n0.0 1.0 0.0\n0.0 0.0 1.0\nend unit_cell_cart\n\nbegin atoms_cart\nFe 0.0 0.0 0.0\nend atoms_cart\n\nbegin projections\nFe:s\nFe:p\nend projections\n";
        let hr = "generated\n4\n1\n1\n0 0 0 1 1 0.0 0.0\n0 0 0 1 2 0.0 0.0\n0 0 0 1 3 0.0 0.0\n0 0 0 1 4 0.0 0.0\n0 0 0 2 1 0.0 0.0\n0 0 0 2 2 0.0 0.0\n0 0 0 2 3 0.0 0.0\n0 0 0 2 4 0.0 0.0\n0 0 0 3 1 0.0 0.0\n0 0 0 3 2 0.0 0.0\n0 0 0 3 3 0.0 0.0\n0 0 0 3 4 0.0 0.0\n0 0 0 4 1 0.0 0.0\n0 0 0 4 2 0.0 0.0\n0 0 0 4 3 0.0 0.0\n0 0 0 4 4 0.0 0.0\n";
        let mut f = fs::File::create(format!("{dir}seedname.win")).unwrap();
        f.write_all(win.as_bytes()).unwrap();
        let mut f = fs::File::create(format!("{dir}seedname_hr.dat")).unwrap();
        f.write_all(hr.as_bytes()).unwrap();
        let model = <Model<false, 3> as Wannier90>::from_hr(dir, "seedname", 0.0).unwrap();
        assert_eq!(model.norb(), 4);
        assert_eq!(
            model.natom(),
            1,
            "multi-line projections must merge into one Atom"
        );
        assert_eq!(model.atoms[0].norb(), 4);
        fs::remove_dir_all(dir).ok();
    }
}