# -*- coding: utf-8 -*-
"""
write-arrow_phonopy.py
A script to extract phonon mode eigenvectors from a Phonopy `band.yaml` file
and convert them into VESTA visualization files (.vesta) with displacement arrows.

Usage:
    python write-arrow_phonopy.py [band.yaml] [template.vesta] [options]

Inputs:
    1. band.yaml: Phonopy calculation result file. 
       (Note: Must be generated with the `--eigvecs` or `--eigenvectors` flag or `EIGENVECTORS = .TRUE.` in Phonopy)
    2. template.vesta: A VESTA file containing the structure to which arrows will be added.

Outputs:
    - Multiple VESTA files, each representing a phonon mode with vibration arrows.
    - Filename format: ARROW-QPOINTS=[coords]-MODE[index]-freq=[frequency]cm-1.vesta

Key Options:
    -s, --scale: Scaling factor for the arrow lengths (Default: 5.0).
    -qg, --qpoint-gamma: Process only the Gamma point (0, 0, 0).
    -qs, --qpoint-special_point: Process only high-symmetry points defined in band.yaml.
    -q, --qpoint [index]: Process only a specific q-point by its index (starting from 0).
"""

import numpy as np
import sys
import os
import shutil
import argparse

def get_args():
    parser = argparse.ArgumentParser(description="Extract Phonon Vectors to VESTA files from Phonopy band.yaml")
    
    # Files
    parser.add_argument('band_yaml', nargs='?', default='band.yaml', help='Phonopy band.yaml file (default: band.yaml)')
    parser.add_argument('vesta_file', nargs='?', default='input.vesta', help='VESTA file (default: input.vesta)')
    
    # Parameters
    parser.add_argument('--scale', '-s', type=float, default=5.0, help='Scaling factor for vectors (default: 5.0)')
    
    # Filtering
    parser.add_argument('--qpoint-gamma', '-qg', action='store_true', help='Extract only the Gamma point (0, 0, 0)')
    parser.add_argument('--qpoint-special_point', '-qs', action='store_true', help='Extract only special points (High Symmetry Points) without duplicates')
    parser.add_argument('--qpoint', '-q', type=int, help='Extract a specific q-point by index (0-based)')
    
    args = parser.parse_args()
    return args

def read_files(args):
    band_fn = args.band_yaml
    vesta_fn = args.vesta_file
        
    if not os.path.exists(band_fn):
        print(f"ERROR: Phonopy band file '{band_fn}' not found.")
        print(f"Try: python {sys.argv[0]} your_band.yaml {vesta_fn}")
        sys.exit(1)
        
    if not os.path.exists(vesta_fn):
        print(f"ERROR: VESTA input file '{vesta_fn}' not found.")
        print(f"Please provide a VESTA file as the second argument, or ensure '{vesta_fn}' exists in the current directory.")
        print(f"Usage example: python {sys.argv[0]} {band_fn} POSCAR.vesta")
        sys.exit(1)

    try:
        content = open(band_fn, 'r').read()
        if 'phonon:' not in content:
            print(f"ERROR: '{band_fn}' does not appear to be a valid Phonopy band.yaml file (missing 'phonon:' tag).")
            sys.exit(1)
            
        parts = content.split('phonon:')
        header = parts[0]
        phonon_data = parts[1]
        
        vesta = open(vesta_fn, 'r').read()
    except Exception as e:
        print(f"Error reading files: {e}")
        sys.exit(1)

    print(f"Reading from: {band_fn} and {vesta_fn}")
    return header, phonon_data, vesta

def parse_header(header):
    lines = header.split('\n')
    segments = []
    labels = []
    in_segment = False
    in_labels = False
    
    for line in lines:
        sline = line.strip()
        if sline.startswith('segment_nqpoint:'):
            in_segment = True
            in_labels = False
            continue
        elif sline.startswith('labels:'):
            in_labels = True
            in_segment = False
            continue
        elif sline.endswith(':') and not sline.startswith('-'):
            in_segment = False
            in_labels = False

        if in_segment and sline.startswith('-'):
            try:
                segments.append(int(sline.replace('-', '').strip()))
            except ValueError: pass
                    
        if in_labels and sline.startswith('-'):
            try:
                content = sline.split('[')[1].split(']')[0]
                parts = [p.strip().replace("'", "").replace('"', "") for p in content.split(',')]
                if len(parts) >= 2: labels.append((parts[0], parts[1]))
            except: pass
    return segments, labels

def get_special_points_map(segments, labels):
    if not segments or not labels: return {}
    special_points = {}
    current_idx = 0
    for i, seg_len in enumerate(segments):
        start_idx = current_idx
        end_idx = current_idx + seg_len - 1
        if i < len(labels):
            start_label, end_label = labels[i][0], labels[i][1]
        else:
            start_label, end_label = "SP", "SP"
        special_points[start_idx] = start_label
        if i == len(segments) - 1:
            special_points[end_idx] = end_label
        current_idx += seg_len
    return special_points

def filter_qpoints(phonon_data, header, args):
    raw_blocks = phonon_data.split('q-position:')
    q_blocks = raw_blocks[1:]
    segments, labels = parse_header(header)
    special_points_map = get_special_points_map(segments, labels)
    
    filtered_blocks = []
    filtered_labels = []
    
    if not (args.qpoint_gamma or args.qpoint_special_point or args.qpoint is not None):
        return q_blocks, [None] * len(q_blocks)

    for idx, block in enumerate(q_blocks):
        keep = False
        label = None
        if args.qpoint is not None and idx == args.qpoint:
            keep = True
        if args.qpoint_special_point and idx in special_points_map:
            keep = True
            label = special_points_map[idx]
        if args.qpoint_gamma:
            try:
                pos_str = block.split('[')[1].split(']')[0]
                vec = np.array([float(x) for x in pos_str.split(',')])
                if np.linalg.norm(vec) < 1e-4:
                    keep = True
                    label = "GM" if label is None else label
            except: pass
        if keep:
            filtered_blocks.append(block)
            filtered_labels.append(label)
    return filtered_blocks, filtered_labels

def extract(q_blocks) :
    q_position = []
    bands = []
    for block in q_blocks:
        try:
            q_position.append(block.split('[')[1].split(']')[0])
            band_parts = block.split('frequency:')[1:]
            if band_parts: bands.append(band_parts)
        except IndexError: continue
            
    if not bands:
        print("No valid band data found.")
        sys.exit(0)

    nq_point = len(bands)
    nbands = len(bands[0])
    qpoint_band = np.zeros((nq_point, nbands))
    eigenvectors = [[None for b in range(nbands)] for q in range(nq_point)]
    
    for q_point in range(nq_point) :
        for band in range(nbands) :
            content = bands[q_point][band]
            qpoint_band[q_point][band] = float(content.strip().split()[0])
            if 'atom' in content:
                eigenvectors[q_point][band] = content.split('atom')[1:]
            else:
                print(f"Warning: No eigenvectors found for Q-point {q_point}, Band {band+1}.")
                sys.exit(1)

    natoms = len(eigenvectors[0][0]) 
    displacements = np.zeros((nq_point, nbands, natoms, 3))
    
    for q_point in range(nq_point) :
        for band in range(nbands) :
            for atom in range(natoms) :
                vector = eigenvectors[q_point][band][atom]
                for direction in range(3) :
                    data = float(vector.split('[')[direction + 1].split(',')[0])
                    displacements[q_point][band][atom][direction] = data
    return(displacements, qpoint_band, q_position)
  
def write_VESTA(vesta, displacements, qpoint_band, q_position, filtered_labels, scale):
    nqpoint, nbands, natoms, _ = displacements.shape
    
    # Unit conversion: 1 THz = 33.35640951981521 cm-1
    THZ_TO_CM1 = 33.35640951981521
    
    print(f"Writing VESTA files to current directory...")
    
    for qpoint in range(nqpoint) :
        try:
            coords = [float(x) for x in q_position[qpoint].split(',')]
            # Check for Gamma point
            if np.linalg.norm(coords) < 1e-4:
                q_pos_str = "GAMMA"
            else:
                # Format: 0.33_0.00_0.00
                q_pos_str = '%.2f_%.2f_%.2f' % (coords[0], coords[1], coords[2])
        except:
            q_pos_str = f"q_{qpoint}"

        for band in range(nbands) :
            freq_thz = qpoint_band[qpoint][band]
            freq_cm1 = freq_thz * THZ_TO_CM1
            
            # Filename: ARROW-QPOINTS=[qpoints]-MODE[mode]-freq=[frequency]cm-1.vesta
            filename = "ARROW-QPOINTS=%s-MODE%03d-freq=%.2fcm-1.vesta" % (q_pos_str, band + 1, freq_cm1)
            
            towrite = vesta.split('VECTR')[0] + 'VECTR\n'
            for atom in range(natoms) :
                towrite += '%5d%10.5f%10.5f%10.5f\n' % (atom + 1, displacements[qpoint][band][atom][0] * scale, displacements[qpoint][band][atom][1] * scale, displacements[qpoint][band][atom][2] * scale)
                towrite += '%5d 0 0 0 0\n  0 0 0 0 0\n' % (atom + 1)
            towrite += '0 0 0 0 0\nVECTT\n'
            for atom in range(natoms) :
                towrite += '%5d  0.500 255   0   0 0\n' % (atom + 1)
            towrite += '0 0 0 0 0\nSPLAN' + vesta.split('SPLAN')[1]
            
            with open(filename, 'w') as f: f.write(towrite)
            
    print(f"Done.")
    return(0)
    
if __name__ == "__main__":
    args = get_args()
    header, phonon_data, vesta_content = read_files(args)
    filtered_q_blocks, filtered_labels = filter_qpoints(phonon_data, header, args)
    if not filtered_q_blocks:
        print("No q-points matched the criteria."); sys.exit(0)
    displacements, qpoint_band, q_position = extract(filtered_q_blocks)
    write_VESTA(vesta_content, displacements, qpoint_band, q_position, filtered_labels, args.scale)
