# Python script to generate Quantum ESPRESSO input files for GaAs
# using pymatgen for CIF parsing and k-point path generation

from pymatgen.core import Structure, Element
from pymatgen.symmetry.bandstructure import HighSymmKpath

# Read the structure from the CIF file
structure = Structure.from_file("GaAs.cif")

# Extract unique atomic species in the structure (preserve order of appearance)
species_in_structure = []
for site in structure:
    symbol = site.specie.symbol  # e.g., "Ga" or "As"
    if symbol not in species_in_structure:
        species_in_structure.append(symbol)

# Prepare ATOMIC_SPECIES lines with atomic mass and pseudopotential filenames
atomic_species_lines = []
for symbol in species_in_structure:
    element = Element(symbol)
    mass = element.atomic_mass  # atomic mass in amu
    # Convert to float if needed (Element.atomic_mass might be a FloatWithUnit)
    mass_val = float(mass) if hasattr(mass, "__float__") else mass
    pseudo_file = f"{symbol}.pbe-n-van.UPF"
    atomic_species_lines.append(f"  {symbol} {mass_val:.6f} {pseudo_file}")

# Prepare ATOMIC_POSITIONS lines in fractional (crystal) coordinates
atomic_positions_lines = []
for site in structure:
    sym = site.specie.symbol
    x, y, z = site.frac_coords
    atomic_positions_lines.append(f"  {sym} {x:.6f} {y:.6f} {z:.6f}")

# Prepare CELL_PARAMETERS lines in Angstrom units
cell_parameters_lines = []
for vec in structure.lattice.matrix:  # 3x3 matrix of lattice vectors
    cell_parameters_lines.append(f"  {vec[0]:.6f} {vec[1]:.6f} {vec[2]:.6f}")

# Common input sections for all calculations
control_block = """&CONTROL
  prefix = 'GaAs',
  pseudo_dir = './',
  outdir = './tmp',
  verbosity = 'high',
  calculation = '{calc}'
/
"""
system_block = f"""&SYSTEM
  ibrav = 0,
  nat = {structure.num_sites},
  ntyp = {len(species_in_structure)},
  ecutwfc = 30,
  ecutrho = 240,
  occupations = 'smearing',
  smearing = 'marzari-vanderbilt',
  degauss = 0.02{',\\n  nbnd = {nbnd}' if False else ''}
/
"""
# (Note: nbnd will be added later for the bands calculation)
electrons_block = """&ELECTRONS
  conv_thr = 1.0d-8,
  mixing_beta = 0.7
/
"""

# Build SCF input content
scf_input = control_block.format(calc="scf")
scf_input += system_block.format(nbnd="")  # no nbnd for SCF
scf_input += electrons_block
scf_input += "ATOMIC_SPECIES\n" + "\n".join(atomic_species_lines) + "\n"
scf_input += "ATOMIC_POSITIONS (crystal)\n" + "\n".join(atomic_positions_lines) + "\n"
scf_input += "CELL_PARAMETERS (angstrom)\n" + "\n".join(cell_parameters_lines) + "\n"
scf_input += "K_POINTS {automatic}\n6 6 6 0 0 0\n"  # 6x6x6 k-point mesh (Gamma-centered)

# Build NSCF (DOS) input content
nscf_input = control_block.format(calc="nscf")
nscf_input += system_block.format(nbnd="")  # no nbnd for NSCF DOS
nscf_input += electrons_block
nscf_input += "ATOMIC_SPECIES\n" + "\n".join(atomic_species_lines) + "\n"
nscf_input += "ATOMIC_POSITIONS (crystal)\n" + "\n".join(atomic_positions_lines) + "\n"
nscf_input += "CELL_PARAMETERS (angstrom)\n" + "\n".join(cell_parameters_lines) + "\n"
nscf_input += "K_POINTS {automatic}\n12 12 12 0 0 0\n"  # 12x12x12 k-point mesh for DOS

# Build bands calculation input content
bands_input = control_block.format(calc="bands")
# Estimate a reasonable number of bands (nbnd) for band structure (e.g., double the valence band count)
nbnd_val = 12  # example value; adjust if needed for more unoccupied bands
bands_input += f"""&SYSTEM
  ibrav = 0,
  nat = {structure.num_sites},
  ntyp = {len(species_in_structure)},
  ecutwfc = 30,
  ecutrho = 240,
  occupations = 'smearing',
  smearing = 'marzari-vanderbilt',
  degauss = 0.02,
  nbnd = {nbnd_val}
/
"""
bands_input += electrons_block

# Generate high-symmetry k-point path for band structure using pymatgen
kpath = HighSymmKpath(structure)
kpath_data = kpath.kpath
segments = kpath_data["path"]           # list of k-point label segments, e.g. [('Γ','X'), ('X','L'), ...]
kpoints_dict = kpath_data["kpoints"]    # dict of fractional coordinates for special k-points

# Prepare K_POINTS lines in 'crystal_b' format for the band path
band_kpoints_lines = []
if segments:
    # Start with the first point of the first segment
    start_label = segments[0][0]
    start_k = kpoints_dict[start_label]
    # Replace Greek Gamma symbol with 'Gamma' for readability in comments
    label_comment = start_label.replace("Γ", "Gamma")
    # Placeholder 0 for subdivisions (to be replaced when next point is added)
    band_kpoints_lines.append(f"  {start_k[0]:.6f} {start_k[1]:.6f} {start_k[2]:.6f} 0  ! {label_comment}")
    # Iterate over each segment and add the end-point of the segment
    for seg in segments:
        end_label = seg[-1]
        end_k = kpoints_dict[end_label]
        label_comment = end_label.replace("Γ", "Gamma")
        # Set subdivisions for the previous line (number of points between this end point and the previous point)
        if band_kpoints_lines:
            subdivisions = 20  # default subdivisions for each segment
            band_kpoints_lines[-1] = band_kpoints_lines[-1].replace(" 0 ", f" {subdivisions} ")
        # Add current end-point with a placeholder (will be replaced if another segment follows)
        band_kpoints_lines.append(f"  {end_k[0]:.6f} {end_k[1]:.6f} {end_k[2]:.6f} 0  ! {label_comment}")
    # For the last k-point, the subdivision number is not used. We leave it as 0.
else:
    # In case no path was found (should not happen for periodic crystals), use Γ as a single point
    band_kpoints_lines.append("  0.000000 0.000000 0.000000 1  ! Gamma")

bands_input += "ATOMIC_SPECIES\n" + "\n".join(atomic_species_lines) + "\n"
bands_input += "ATOMIC_POSITIONS (crystal)\n" + "\n".join(atomic_positions_lines) + "\n"
bands_input += "CELL_PARAMETERS (angstrom)\n" + "\n".join(cell_parameters_lines) + "\n"
bands_input += "K_POINTS {crystal_b}\n"
bands_input += f"{len(band_kpoints_lines)}\n" + "\n".join(band_kpoints_lines) + "\n"

# Write the input strings to files
with open("GaAs_scf.in", "w", newline="\n") as f:
    f.write(scf_input)
with open("GaAs_nscf_dos.in", "w", newline="\n") as f:
    f.write(nscf_input)
with open("GaAs_bands.in", "w", newline="\n") as f:
    f.write(bands_input)

print("Generated GaAs_scf.in, GaAs_nscf_dos.in, and GaAs_bands.in")