"""Test basic loading and distance calculation for EPW data.

PURPOSE:
    Verify the distance_resolved_couplings methods work without hanging.

USAGE:
    uv run python agent_files/debug_epw_distance/test_basic.py
"""

import os
import sys
import time

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

from HamiltonIO.epw.epwparser import Epmat, read_crystal_fmt
from ase.units import Bohr


def main():
    # Find test data
    base = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
    data_dir = os.path.join(base, "HamiltonIO", "epw", "test", "up")

    print(f"Loading data from: {data_dir}")

    # Read crystal structure
    t0 = time.time()
    crystal = read_crystal_fmt(os.path.join(data_dir, "crystal.fmt"))
    print(f"Crystal loaded in {time.time() - t0:.2f}s")
    print(f"  natom={crystal.natom}, nmode={crystal.nmode}")

    # Read EPW data
    t0 = time.time()
    ep = Epmat()
    ep.crystal = crystal
    ep.read(path=data_dir, prefix="SrMnO3", epmat_ncfile="epmat.nc")
    print(f"Epmat loaded in {time.time() - t0:.2f}s")
    if hasattr(ep, "epmatfile") and ep.epmatfile is not None:
        print(f"  nwann={ep.nwann}, using netCDF file")
    elif hasattr(ep, "epmat_wann") and ep.epmat_wann is not None:
        print(f"  nwann={ep.nwann}, shape={ep.epmat_wann.shape}")
    else:
        print(f"  nwann={ep.nwann}, no matrix data loaded!")

    # Get cell
    cell = crystal.at.reshape(3, 3) * crystal.alat * Bohr
    print(f"Cell (Angstrom):\n{cell}")

    # Test Rk-based distance calculation
    imode = 0
    print(f"\nTesting Rk-based distances for mode {imode}...")
    t0 = time.time()
    entries_Rk = ep.distance_resolved_couplings_Rk(imode=imode, cell=cell)
    print(f"Rk calculation done in {time.time() - t0:.2f}s")
    print(f"  Found {len(entries_Rk)} entries")
    if entries_Rk:
        print(f"  Sample entry: {entries_Rk[0]}")
        # Print distance and magnitude statistics
        distances = [e["distance"] for e in entries_Rk]
        mags = [abs(e["coupling"]) for e in entries_Rk]
        print(f"  Distance range: {min(distances):.2f} - {max(distances):.2f} Angstrom")
        print(f"  |g| range: {min(mags):.4e} - {max(mags):.4e} eV/Angstrom")

    # Test Rg-based distance calculation
    print(f"\nTesting Rg-based distances for mode {imode}...")
    t0 = time.time()
    entries_Rg = ep.distance_resolved_couplings_Rg(imode=imode, cell=cell)
    print(f"Rg calculation done in {time.time() - t0:.2f}s")
    print(f"  Found {len(entries_Rg)} entries")
    if entries_Rg:
        print(f"  Sample entry: {entries_Rg[0]}")
        distances = [e["distance"] for e in entries_Rg]
        mags = [abs(e["coupling"]) for e in entries_Rg]
        print(f"  Distance range: {min(distances):.2f} - {max(distances):.2f} Angstrom")
        print(f"  |g| range: {min(mags):.4e} - {max(mags):.4e} eV/Angstrom")

    print("\nSuccess!")


if __name__ == "__main__":
    main()
