import numpy as np

def read_chgcar(filename):
    info = {}
    with open(filename, 'r') as f:
        # 1. タイトル
        info['title'] = f.readline().strip()
        
        # 2. スケール因子
        info['scale'] = float(f.readline().strip())
        
        # 3-5. 格子ベクトル
        lattice = []
        for _ in range(3):
            lattice.append([float(x) for x in f.readline().split()])
        info['lattice'] = np.array(lattice)
        
        # 6. 原子種
        info['atoms'] = f.readline().split()
        
        # 7. 原子数
        info['numbers'] = [int(x) for x in f.readline().split()]
        
        # 8. 座標の種類
        l = f.readline()
        info['coord_type'] = l.strip()
        
        # 9. 原子座標
        natoms = sum(info['numbers'])
        coords = []
        for _ in range(natoms):
            l = f.readline()
            coords.append([float(x) for x in l.split()[:3]])
        info['coords'] = np.array(coords)
        
        # 格子分割数
        while True:
            l = f.readline()
            if l.strip() != '':
                break

        nx, ny, nz = [int(x) for x in l.split()]
        info['grid'] = (nx, ny, nz)
        
        # 体積データ読み込み
        ngrid = nx * ny * nz
        data = []
        while len(data) < ngrid:
            line = f.readline()
            if not line:
                break
            data.extend([float(x) for x in line.split()])
        
        if len(data) != ngrid:
            raise ValueError("CHGCAR data size mismatch")
        
        # P[ix][iy][iz] に変換（ネストリスト）
        arr = np.array(data).reshape((nx, ny, nz), order='F')  # Fortran順
        P = [[[arr[i, j, k] for k in range(nz)]
              for j in range(ny)]
              for i in range(nx)]
        
    return info, P

def main(infile):
    info, P = read_chgcar(infile)
    print(info['title'])
    print(info['grid'])
    print(P[0][0][0])  # (ix=0, iy=0, iz=0) の電荷密度

if __name__ == '__main__':
    main('CHGCAR')