from pathlib import Path

import numpy as np
from tqdm import tqdm


if __name__ == '__main__':
    
    all_mu = Path("22-05-13-LPS-no-POTCAR").rglob("mu.txt")

    for path in tqdm(list(all_mu)):
        X = np.loadtxt(path, skiprows=3)
        where_nonzero = np.where(X[:, 1:].sum(axis=1) != 0)[0]
        x0 = max(where_nonzero[0] - 100, 0)
        xf = min(where_nonzero[-1] + 100, 40000)
        
        pf = Path(path.parent) / "mu.dat"
        np.savetxt(pf, X[x0:xf, :4])