#!/usr/bin/env python3
# -*- coding: ascii -*-
"""
Randomly swap one Li-Y pair in a VASP5 POSCAR and write N new files.
- Keeps coordinate mode (Direct/Cartesian), Selective Dynamics flags, and Velocities.
- Requires VASP5+ POSCAR (with element symbols line).
"""

import random
from pathlib import Path



def _is_ints(tokens):
    try:
        _ = [int(t) for t in tokens]
        return True
    except ValueError:
        return False

def _find_species_index(species, target):
    t = target.lower()
    for i, s in enumerate(species):
        if s.lower() == t:
            return i
    raise ValueError("Element '%s' not found in species line: %s" % (target, " ".join(species)))

def parse_poscar(lines):
    if len(lines) < 8:
        raise ValueError("POSCAR too short")

    comment = lines[0].rstrip("\n")
    scale = lines[1].rstrip("\n")
    lv = [lines[2].rstrip("\n"), lines[3].rstrip("\n"), lines[4].rstrip("\n")]

    i = 5
    tokens = lines[i].split()
    if _is_ints(tokens):
        raise ValueError("VASP4 format (no element symbols line) is not supported.")
    species = tokens
    i += 1

    counts_line = lines[i].rstrip("\n")
    counts = [int(x) for x in counts_line.split()]
    i += 1

    selective = False
    sel_line = None
    if i < len(lines) and lines[i].strip().lower().startswith("s"):
        selective = True
        sel_line = lines[i].rstrip("\n")
        i += 1

    coord_line = lines[i].rstrip("\n")
    i += 1

    natoms = sum(counts)
    pos_start = i
    pos_end = i + natoms
    if pos_end > len(lines):
        raise ValueError("Atom count exceeds file length.")
    pos_lines = [lines[j].rstrip("\n") for j in range(pos_start, pos_end)]

    # Optional Velocities section
    k = pos_end
    while k < len(lines) and lines[k].strip() == "":
        k += 1

    vel_header = None
    vel_lines = None
    vel_end = k
    if k < len(lines) and lines[k].strip().lower().startswith("veloc"):
        vel_header = lines[k].rstrip("\n")
        vel_start = k + 1
        vel_end = vel_start + natoms
        if vel_end > len(lines):
            raise ValueError("Velocities section length inconsistent with atom count.")
        vel_lines = [lines[j].rstrip("\n") for j in range(vel_start, vel_end)]

    tail_lines = [lines[j].rstrip("\n") for j in range(vel_end, len(lines))] if vel_end < len(lines) else []

    # Species index ranges (by order in species list)
    starts = []
    acc = 0
    for c in counts:
        starts.append(acc)
        acc += c
    species_ranges = []
    for idx, sp in enumerate(species):
        s, e = starts[idx], starts[idx] + counts[idx]
        species_ranges.append((sp, s, e))

    return {
        "comment": comment,
        "scale": scale,
        "lv": lv,
        "species": species,
        "counts": counts,
        "counts_line": counts_line,
        "selective": selective,
        "sel_line": sel_line,
        "coord_line": coord_line,
        "natoms": natoms,
        "pos_lines": pos_lines,
        "vel_header": vel_header,
        "vel_lines": vel_lines,
        "tail_lines": tail_lines,
        "species_ranges": species_ranges,
    }

def build_poscar(data, pos_lines, vel_lines=None):
    out = []
    out.append(data["comment"])
    out.append(data["scale"])
    out.extend(data["lv"])
    out.append(" ".join(data["species"]))
    out.append(data["counts_line"])
    if data["selective"]:
        out.append(data["sel_line"] if data["sel_line"] is not None else "Selective dynamics")
    out.append(data["coord_line"])
    out.extend(pos_lines)
    if data["vel_header"] is not None and vel_lines is not None:
        out.append(data["vel_header"])
        out.extend(vel_lines)
    if data["tail_lines"]:
        out.extend(data["tail_lines"])
    return "\n".join(out) + "\n"

def _swap_once(data, rng, li_label="Li", y_label="Y"):
    si_li = _find_species_index(data["species"], li_label)
    si_y  = _find_species_index(data["species"], y_label)
    _, li_start, li_end = data["species_ranges"][si_li]
    _, y_start,  y_end  = data["species_ranges"][si_y]

    li_pick = rng.randrange(li_start, li_end)
    y_pick  = rng.randrange(y_start, y_end)

    new_pos = list(data["pos_lines"])
    new_pos[li_pick], new_pos[y_pick] = new_pos[y_pick], new_pos[li_pick]

    new_vel = None
    if data["vel_lines"] is not None:
        new_vel = list(data["vel_lines"])
        new_vel[li_pick], new_vel[y_pick] = new_vel[y_pick], new_vel[li_pick]

    return new_pos, new_vel, (li_pick, y_pick)

def swap(n, input_file, output_dir):
    """
    Generate n POSCAR files, each with one random Li-Y swap.

    Returns: list of Path to written files.
    """
    input_path = Path(input_file)
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    lines = input_path.read_text().splitlines()
    data = parse_poscar(lines)

    rng = random.Random()
    base = input_path.name

    out_paths = []
    for k in range(1, n + 1):
        new_pos, new_vel, picked = _swap_once(data, rng)
        txt = build_poscar(data, new_pos, new_vel)
        out_path = out_dir / f"swap_{k}_{base}"
        out_path.write_text(txt)
        out_paths.append(out_path)
        print(f"Wrote {out_path}  (swapped Li idx {picked[0]} <-> Y idx {picked[1]})")
    return out_paths
# --------- Editable defaults for direct run ---------
INPUT_FILE = "data_POSCAR/origin/p3m1.vasp"   # path to input POSCAR
OUTPUT_DIR = "data_POSCAR/p3m1"    # output directory
N = 5                   # number of files to generate
# ----------------------------------------------------
if __name__ == "__main__":
    # Direct-run entry: edit INPUT_FILE/OUTPUT_DIR/N above to change behavior.
    swap(n=N, input_file=INPUT_FILE, output_dir=OUTPUT_DIR)