#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
sample_tkcif_tkcrystal.py

Sample program:
    1. Read one CIF file through tkcif_reader.read_structure()
       This uses the tkcif-side reading/fallback layer.
    2. Extract crystal-structure information through tkcrystal wrapper functions.
    3. Print the result to console.

Expected files/packages:
    - tkcif_reader.py
    - tkcif_normalize.py
    - tkcrystal/ package

Basic usage:
    python sample_tkcif_tkcrystal.py sample.cif

With XRD and symmetry operations:
    python sample_tkcif_tkcrystal.py sample.cif --xrd 1 --symops 1

JSON output:
    python sample_tkcif_tkcrystal.py sample.cif --mode json

Notes:
    - read_structure() is intentionally imported from tkcif_reader.
    - get_*_inf() functions are imported from tkcrystal.
    - Return values from tkcrystal functions are plain dictionaries.
"""

from __future__ import annotations

import argparse
import json
import sys
import traceback
from pathlib import Path
from pprint import pprint
from typing import Any


def import_tkcif_reader():
    try:
        from tkcif_reader import read_structure
        return read_structure
    except Exception as exc:
        print("Error: failed to import tkcif_reader.read_structure", file=sys.stderr)
        print(f"  {type(exc).__name__}: {exc}", file=sys.stderr)
        print("\nMake sure tkcif_reader.py is in the same directory or PYTHONPATH.", file=sys.stderr)
        raise


def import_tkcrystal_wrappers():
    try:
        from tkcrystal import (
            get_atom_inf,
            get_composition_inf,
            get_crystal_inf,
            get_density_inf,
            get_lattice_inf,
            get_site_inf,
            get_spg_inf,
            get_symmetry_operations_inf,
            get_xrd_inf,
        )

        # printer.py is optional, but useful for readable text output.
        try:
            from tkcrystal.printer import crystal_inf_to_text
        except Exception:
            crystal_inf_to_text = None

        return {
            "get_atom_inf": get_atom_inf,
            "get_composition_inf": get_composition_inf,
            "get_crystal_inf": get_crystal_inf,
            "get_density_inf": get_density_inf,
            "get_lattice_inf": get_lattice_inf,
            "get_site_inf": get_site_inf,
            "get_spg_inf": get_spg_inf,
            "get_symmetry_operations_inf": get_symmetry_operations_inf,
            "get_xrd_inf": get_xrd_inf,
            "crystal_inf_to_text": crystal_inf_to_text,
        }

    except Exception as exc:
        print("Error: failed to import tkcrystal wrapper functions", file=sys.stderr)
        print(f"  {type(exc).__name__}: {exc}", file=sys.stderr)
        print("\nMake sure the tkcrystal package directory is in the same directory or PYTHONPATH.", file=sys.stderr)
        raise


def print_dict(data: dict[str, Any], *, mode: str = "pprint") -> None:
    if mode == "json":
        print(json.dumps(data, ensure_ascii=False, indent=2))
    elif mode == "pprint":
        pprint(data, width=120, sort_dicts=False)
    else:
        raise ValueError(f"Unknown dictionary print mode: {mode}")


def compact_site_table(site_info: dict[str, Any]) -> str:
    lines: list[str] = []
    lines.append("Atomic sites")
    lines.append("  #    species                 frac_x        frac_y        frac_z")
    for site in site_info.get("sites", []):
        fc = site.get("frac_coords", ["", "", ""])
        try:
            x, y, z = float(fc[0]), float(fc[1]), float(fc[2])
            lines.append(
                f"  {int(site.get('index', 0)):4d} "
                f"{str(site.get('species_string', '')):22s} "
                f"{x:12.7f} {y:12.7f} {z:12.7f}"
            )
        except Exception:
            lines.append(f"  {site}")
    return "\n".join(lines)


def print_basic_text(
    *,
    path: Path,
    read_info: Any,
    crystal_info: dict[str, Any],
    lattice_info: dict[str, Any],
    composition_info: dict[str, Any],
    density_info: dict[str, Any],
    spg_info: dict[str, Any],
    site_info: dict[str, Any],
    crystal_inf_to_text: Any = None,
) -> None:
    print("=" * 72)
    print("Input")
    print("=" * 72)
    print(f"CIF file             : {path}")

    print("")
    print("=" * 72)
    print("tkcif read_structure()")
    print("=" * 72)
    try:
        print(read_info.short_report())
    except Exception:
        pprint(read_info)

    print("")
    print("=" * 72)
    print("tkcrystal get_crystal_inf()")
    print("=" * 72)

    if crystal_inf_to_text is not None:
        print(crystal_inf_to_text(crystal_info))
    else:
        print(f"Formula              : {composition_info.get('formula')}")
        print(f"Reduced formula      : {composition_info.get('reduced_formula')}")
        print(f"Number of sites      : {site_info.get('n_sites')}")
        print(f"Volume [A^3]         : {lattice_info.get('volume')}")
        print(f"Density [g/cm^3]     : {density_info.get('mass_density_g_cm3')}")
        print("")
        print("Lattice")
        print(
            "  a, b, c [A]        : "
            f"{lattice_info.get('a')}, {lattice_info.get('b')}, {lattice_info.get('c')}"
        )
        print(
            "  alpha,beta,gamma   : "
            f"{lattice_info.get('alpha')}, {lattice_info.get('beta')}, {lattice_info.get('gamma')}"
        )
        print("")
        print("Space group")
        print(f"  Symbol             : {spg_info.get('symbol')}")
        print(f"  Number             : {spg_info.get('number')}")
        print(f"  Crystal system     : {spg_info.get('crystal_system')}")
        print("")
        print(compact_site_table(site_info))


def print_xrd_text(xrd_info: dict[str, Any]) -> None:
    print("")
    print("=" * 72)
    print("XRD peaks")
    print("=" * 72)
    print(f"Wavelength           : {xrd_info.get('wavelength')}")
    print(f"Wavelength [A]       : {xrd_info.get('wavelength_angstrom')}")
    print(f"2theta range         : {xrd_info.get('two_theta_range')}")
    print(f"Number of peaks      : {xrd_info.get('n_peaks')}")
    print("")
    print("  #      2theta       d_hkl       intensity     hkl       mult")
    for peak in xrd_info.get("peaks", []):
        hkl = peak.get("hkl")
        hkl_str = " ".join(str(x) for x in hkl) if hkl is not None else ""
        print(
            f"  {int(peak.get('index', 0)):4d} "
            f"{float(peak.get('two_theta', 0.0)):11.5f} "
            f"{float(peak.get('d_hkl', 0.0)):10.5f} "
            f"{float(peak.get('intensity', 0.0)):12.5f} "
            f"{hkl_str:9s} "
            f"{str(peak.get('multiplicity')):>5s}"
        )


def print_symops_text(symops_info: dict[str, Any]) -> None:
    print("")
    print("=" * 72)
    print("Symmetry operations")
    print("=" * 72)
    print(f"Number of operations : {symops_info.get('n_operations')}")
    print(f"Returned operations  : {symops_info.get('n_returned_operations')}")
    print("")
    for op in symops_info.get("operations", []):
        print(f"  #{int(op.get('index', 0)):03d} {op.get('kind', ''):28s} {op.get('xyz', '')}")


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Read CIF by tkcif_reader and print crystal information by tkcrystal wrappers."
    )
    parser.add_argument("cif_file", type=str, help="Input CIF file")
    parser.add_argument(
        "--backend",
        type=str,
        default="auto",
        help="Backend for tkcrystal get_*_inf() functions. Default: auto",
    )
    parser.add_argument(
        "--primitive",
        type=int,
        default=0,
        choices=[0, 1],
        help="Pass primitive=0/1 to tkcif_reader.read_structure(). Default: 0",
    )
    parser.add_argument(
        "--symprec",
        type=float,
        default=1.0e-3,
        help="Symmetry tolerance for get_spg_inf(). Default: 1e-3",
    )
    parser.add_argument(
        "--max-sites",
        type=int,
        default=20,
        help="Maximum number of sites to print. Default: 20",
    )
    parser.add_argument(
        "--mode",
        type=str,
        default="text",
        choices=["text", "pprint", "json"],
        help="Output mode. Default: text",
    )
    parser.add_argument(
        "--xrd",
        type=int,
        default=0,
        choices=[0, 1],
        help="Also calculate XRD peak list. Default: 0",
    )
    parser.add_argument(
        "--xray-source",
        type=str,
        default="CuKa1",
        help="X-ray source name or wavelength for XRD. Default: CuKa1",
    )
    parser.add_argument(
        "--twotheta-min",
        type=float,
        default=10.0,
        help="Minimum 2theta for XRD. Default: 10",
    )
    parser.add_argument(
        "--twotheta-max",
        type=float,
        default=80.0,
        help="Maximum 2theta for XRD. Default: 80",
    )
    parser.add_argument(
        "--max-peaks",
        type=int,
        default=30,
        help="Maximum number of XRD peaks to print. Default: 30",
    )
    parser.add_argument(
        "--symops",
        type=int,
        default=0,
        choices=[0, 1],
        help="Also print symmetry operations. Default: 0",
    )
    parser.add_argument(
        "--max-ops",
        type=int,
        default=20,
        help="Maximum number of symmetry operations to print. Default: 20",
    )
    parser.add_argument(
        "--show-traceback",
        type=int,
        default=0,
        choices=[0, 1],
        help="Show traceback on error. Default: 0",
    )
    parser.add_argument(
        "--pause",
        type=int,
        default=0,
        choices=[0, 1],
        help="Pause before exit. Default: 0",
    )

    args = parser.parse_args()

    try:
        path = Path(args.cif_file)
        if not path.is_file():
            print(f"Error: file not found: {path}", file=sys.stderr)
            return 1

        read_structure = import_tkcif_reader()
        wrappers = import_tkcrystal_wrappers()

        structure, read_info = read_structure(
            path,
            primitive=bool(args.primitive),
            return_info=True,
        )

        crystal_info = wrappers["get_crystal_inf"](
            structure,
            backend=args.backend,
            include_sites=True,
            max_sites=args.max_sites,
            include_symmetry=True,
            symprec=args.symprec,
        )

        lattice_info = wrappers["get_lattice_inf"](structure, backend=args.backend)
        composition_info = wrappers["get_composition_inf"](structure, backend=args.backend)
        density_info = wrappers["get_density_inf"](structure, backend=args.backend)
        site_info = wrappers["get_site_inf"](
            structure,
            backend=args.backend,
            max_sites=args.max_sites,
            include_cartesian=True,
        )
        spg_info = wrappers["get_spg_inf"](
            structure,
            backend=args.backend,
            symprec=args.symprec,
        )

        result: dict[str, Any] = {
            "input": {
                "path": str(path),
            },
            "read_info": {
                "backend": getattr(read_info, "backend", None),
                "normalized": getattr(read_info, "normalized", None),
                "encoding": getattr(read_info, "encoding", None),
                "warnings": getattr(read_info, "warnings", []),
                "errors": getattr(read_info, "errors", []),
                "unavailable_backends": getattr(read_info, "unavailable_backends", []),
            },
            "crystal": crystal_info,
            "lattice": lattice_info,
            "composition": composition_info,
            "density": density_info,
            "sites": site_info,
            "spacegroup": spg_info,
        }

        if args.xrd:
            xrd_info = wrappers["get_xrd_inf"](
                structure,
                backend=args.backend,
                wavelength=args.xray_source,
                two_theta_range=(args.twotheta_min, args.twotheta_max),
                max_peaks=args.max_peaks,
            )
            result["xrd"] = xrd_info

        if args.symops:
            symops_info = wrappers["get_symmetry_operations_inf"](
                structure,
                backend=args.backend,
                symprec=args.symprec,
                max_ops=args.max_ops,
            )
            result["symmetry_operations"] = symops_info

        if args.mode == "text":
            print_basic_text(
                path=path,
                read_info=read_info,
                crystal_info=crystal_info,
                lattice_info=lattice_info,
                composition_info=composition_info,
                density_info=density_info,
                spg_info=spg_info,
                site_info=site_info,
                crystal_inf_to_text=wrappers["crystal_inf_to_text"],
            )
            if args.xrd:
                print_xrd_text(result["xrd"])
            if args.symops:
                print_symops_text(result["symmetry_operations"])

        elif args.mode == "pprint":
            print_dict(result, mode="pprint")

        elif args.mode == "json":
            print_dict(result, mode="json")

        return 0

    except Exception as exc:
        print("Error: sample_tkcif_tkcrystal failed", file=sys.stderr)
        print(f"  {type(exc).__name__}: {exc}", file=sys.stderr)
        if args.show_traceback:
            traceback.print_exc()
        return 1

    finally:
        if args.pause:
            input("\nPress ENTER to terminate>>\n")


if __name__ == "__main__":
    raise SystemExit(main())