#!/usr/bin/env python3
import h5py
import numpy as np
import argparse
import sys
import os

def main():
    parser = argparse.ArgumentParser(description="Extract Gruneisen parameters from HDF5 to text format.")
    parser.add_argument("input", nargs='?', default="gruneisen.hdf5", help="Input HDF5 file (default: gruneisen.hdf5)")
    parser.add_argument("-o", "--output", default="gruneisen_full.dat", help="Output data file (default: gruneisen_full.dat)")
    
    args = parser.parse_args()

    input_file = args.input
    output_file = args.output

    if not os.path.isfile(input_file):
        print(f"Error: Input file '{input_file}' not found.")
        sys.exit(1)

    print(f"Reading from {input_file}...")

    try:
        with h5py.File(input_file, 'r') as f:
            # Check required keys
            required_keys = ['qpoint', 'weight', 'frequency', 'gruneisen']
            for key in required_keys:
                if key not in f:
                    print(f"Error: Key '{key}' not found in HDF5 file.")
                    sys.exit(1)

            # Read data
            qpoints = f['qpoint'][:]     # (N_q, 3)
            weights = f['weight'][:]     # (N_q,)
            frequencies = f['frequency'][:] # (N_q, N_band)
            gruneisen = f['gruneisen'][:]   # (N_q, N_band)

            n_q, n_band = frequencies.shape
            
            print(f"Found {n_q} q-points and {n_band} bands.")
            print(f"Writing to {output_file}...")

            with open(output_file, 'w') as out:
                out.write("# qx       qy       qz       weight band_index frequency(THz)  gruneisen_param\n")
                for i in range(n_q):
                    qx, qy, qz = qpoints[i]
                    w = weights[i]
                    for j in range(n_band):
                        freq = frequencies[i, j]
                        gamma = gruneisen[i, j]
                        out.write(f"{qx:10.6f} {qy:10.6f} {qz:10.6f} "
                                  f"{w:6d} {j+1:4d} {freq:15.8f} {gamma:15.8f}\n")
            
            print("Done.")

    except Exception as e:
        print(f"An error occurred: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()
