#!/usr/bin/env python3
"""Convert LAMMPS dump with xs/ys/zs (scaled coords) into x/y/z (absolute coords).

CLI name: rng-xs2x

Supports:
- Orthorhombic boxes: BOX BOUNDS lines are `lo hi` (2 floats)
- Triclinic boxes: BOX BOUNDS lines include tilt factors (3 floats) and the header contains `xy xz yz`

Assumptions:
- ATOMS header includes xs/ys/zs

Streaming conversion (no full-file load).
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass


def die(msg: str) -> None:
    raise SystemExit(f"ERROR: {msg}")


@dataclass
class Box:
    # True box parameters (LAMMPS triclinic convention)
    xlo: float
    xhi: float
    ylo: float
    yhi: float
    zlo: float
    zhi: float
    xy: float
    xz: float
    yz: float

    @property
    def lx(self) -> float:
        return self.xhi - self.xlo

    @property
    def ly(self) -> float:
        return self.yhi - self.ylo

    @property
    def lz(self) -> float:
        return self.zhi - self.zlo


def _bounds_to_true_box(
    *,
    xlo_b: float,
    xhi_b: float,
    ylo_b: float,
    yhi_b: float,
    zlo_b: float,
    zhi_b: float,
    xy: float,
    xz: float,
    yz: float,
) -> Box:
    """Convert LAMMPS triclinic dump bounds to true box corners.

    In a LAMMPS dump with triclinic boxes, the 3 bounds lines are:
      xlo_bound xhi_bound xy
      ylo_bound yhi_bound xz
      zlo_bound zhi_bound yz

    LAMMPS stores the bounds expanded by tilt; recover true xlo/xhi/ylo/yhi via:
      xlo = xlo_b - min(0, xy, xz, xy+xz)
      xhi = xhi_b - max(0, xy, xz, xy+xz)
      ylo = ylo_b - min(0, yz)
      yhi = yhi_b - max(0, yz)
      zlo = zlo_b
      zhi = zhi_b

    Ref: LAMMPS docs (triclinic simulation boxes).
    """

    xlo = xlo_b - min(0.0, xy, xz, xy + xz)
    xhi = xhi_b - max(0.0, xy, xz, xy + xz)
    ylo = ylo_b - min(0.0, yz)
    yhi = yhi_b - max(0.0, yz)
    zlo = zlo_b
    zhi = zhi_b
    return Box(xlo=xlo, xhi=xhi, ylo=ylo, yhi=yhi, zlo=zlo, zhi=zhi, xy=xy, xz=xz, yz=yz)


def _parse_box_bounds(header_line: str, fin) -> tuple[Box, list[str]]:
    """Parse BOX BOUNDS block and return (Box, raw_lines_for_output).

    We preserve raw bounds lines verbatim in output; only coordinates are rewritten.
    """

    # Detect triclinic from header tokens.
    toks = header_line.strip().split()
    triclinic = any(t in ("xy", "xz", "yz") for t in toks)

    raw: list[str] = []
    parts_list: list[list[str]] = []
    for _ in range(3):
        bline = fin.readline()
        if not bline:
            die("EOF while reading BOX BOUNDS")
        raw.append(bline)
        p = bline.split()
        if len(p) < 2:
            die(f"Bad BOX BOUNDS line: {bline.strip()!r}")
        parts_list.append(p)
        if len(p) >= 3:
            triclinic = True

    if not triclinic:
        try:
            xlo, xhi = float(parts_list[0][0]), float(parts_list[0][1])
            ylo, yhi = float(parts_list[1][0]), float(parts_list[1][1])
            zlo, zhi = float(parts_list[2][0]), float(parts_list[2][1])
        except Exception:
            die("Failed to parse orthorhombic BOX BOUNDS")
        box = Box(xlo=xlo, xhi=xhi, ylo=ylo, yhi=yhi, zlo=zlo, zhi=zhi, xy=0.0, xz=0.0, yz=0.0)
        return box, raw

    # Triclinic: expect 3 floats per line.
    if any(len(p) < 3 for p in parts_list):
        die(
            "Triclinic BOX BOUNDS detected (xy/xz/yz), but bounds lines do not have 3 columns. "
            "Expected: 'lo hi tilt' per line."
        )

    try:
        xlo_b, xhi_b, xy = float(parts_list[0][0]), float(parts_list[0][1]), float(parts_list[0][2])
        ylo_b, yhi_b, xz = float(parts_list[1][0]), float(parts_list[1][1]), float(parts_list[1][2])
        zlo_b, zhi_b, yz = float(parts_list[2][0]), float(parts_list[2][1]), float(parts_list[2][2])
    except Exception:
        die("Failed to parse triclinic BOX BOUNDS (lo hi tilt)")

    box = _bounds_to_true_box(
        xlo_b=xlo_b,
        xhi_b=xhi_b,
        ylo_b=ylo_b,
        yhi_b=yhi_b,
        zlo_b=zlo_b,
        zhi_b=zhi_b,
        xy=xy,
        xz=xz,
        yz=yz,
    )
    return box, raw


def _scaled_to_xyz(*, box: Box, xs: float, ys: float, zs: float) -> tuple[float, float, float]:
    """Convert scaled (lamda) coords to xyz.

    For triclinic, LAMMPS uses the h-matrix:
      a = (lx,  0,  0)
      b = (xy, ly,  0)
      c = (xz, yz, lz)

    With lamda = (xs, ys, zs):
      x = xlo + xs*lx + ys*xy + zs*xz
      y = ylo + ys*ly + zs*yz
      z = zlo + zs*lz

    For orthorhombic, xy=xz=yz=0 and this reduces to xlo+xs*lx etc.
    """

    x = box.xlo + xs * box.lx + ys * box.xy + zs * box.xz
    y = box.ylo + ys * box.ly + zs * box.yz
    z = box.zlo + zs * box.lz
    return x, y, z


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(prog="rng-xs2x")
    ap.add_argument("-i", "--input", required=True, help="Input .lammpstrj (with xs/ys/zs)")
    ap.add_argument("-o", "--output", required=True, help="Output .lammpstrj (with x/y/z)")
    ap.add_argument("--precision", type=int, default=8)
    args = ap.parse_args(argv)

    fmt = "{:." + str(args.precision) + "f}"

    frames = 0
    with open(args.input, "r", errors="replace") as fin, open(args.output, "w") as fout:
        while True:
            line = fin.readline()
            if not line:
                break
            if not line.startswith("ITEM: TIMESTEP"):
                die(f"Expected 'ITEM: TIMESTEP', got: {line.strip()!r}")

            fout.write(line)
            timestep = fin.readline()
            if not timestep:
                die("EOF after ITEM: TIMESTEP")
            fout.write(timestep)

            line = fin.readline()
            if not line or not line.startswith("ITEM: NUMBER OF ATOMS"):
                die(f"Expected 'ITEM: NUMBER OF ATOMS', got: {line.strip() if line else 'EOF'}")
            fout.write(line)
            natoms_line = fin.readline()
            if not natoms_line:
                die("EOF after ITEM: NUMBER OF ATOMS")
            fout.write(natoms_line)
            try:
                natoms = int(natoms_line.strip())
            except Exception:
                die(f"Cannot parse atom count: {natoms_line.strip()!r}")

            line = fin.readline()
            if not line or not line.startswith("ITEM: BOX BOUNDS"):
                die(f"Expected 'ITEM: BOX BOUNDS', got: {line.strip() if line else 'EOF'}")
            fout.write(line)

            box, raw_bounds = _parse_box_bounds(line, fin)
            for bline in raw_bounds:
                fout.write(bline)

            line = fin.readline()
            if not line or not line.startswith("ITEM: ATOMS"):
                die(f"Expected 'ITEM: ATOMS ...', got: {line.strip() if line else 'EOF'}")

            header_parts = line.strip().split()
            keys = header_parts[2:]

            try:
                i_xs = keys.index("xs")
                i_ys = keys.index("ys")
                i_zs = keys.index("zs")
            except ValueError:
                die(f"ATOMS header must include xs/ys/zs. Got keys={keys}")

            new_keys = list(keys)
            new_keys[i_xs] = "x"
            new_keys[i_ys] = "y"
            new_keys[i_zs] = "z"
            fout.write("ITEM: ATOMS " + " ".join(new_keys) + "\n")

            for _ in range(natoms):
                aline = fin.readline()
                if not aline:
                    die("EOF while reading ATOMS")
                parts = aline.split()
                if len(parts) < len(keys):
                    die(
                        f"Atom line too short. Expected >= {len(keys)} cols, got {len(parts)}: {aline.strip()!r}"
                    )

                xs = float(parts[i_xs])
                ys = float(parts[i_ys])
                zs = float(parts[i_zs])
                x, y, z = _scaled_to_xyz(box=box, xs=xs, ys=ys, zs=zs)

                parts[i_xs] = fmt.format(x)
                parts[i_ys] = fmt.format(y)
                parts[i_zs] = fmt.format(z)

                fout.write(" ".join(parts) + "\n")

            frames += 1

    print(f"Converted {frames} frame(s): {args.input} -> {args.output}")


if __name__ == "__main__":
    main()
