#!/usr/bin/env python
# infer_symmetry_dof.py

import argparse
import numpy as np

from pymatgen.core import Structure
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.symmetry.groups import SpaceGroup


def wrap01(x):
    return np.array(x, dtype=float) % 1.0


def frac_dist(a, b):
    a = np.array(a, dtype=float)
    b = np.array(b, dtype=float)
    d = (a - b + 0.5) % 1.0 - 0.5
    return np.linalg.norm(d)


def close_frac(a, b, tol=1e-5):
    return frac_dist(a, b) < tol


def close_mod1(a, b, tol=1e-5):
    d = abs((a - b + 0.5) % 1.0 - 0.5)
    return d < tol


def is_fixed_value(v, values=(0.0, 0.25, 1/3, 0.5, 2/3, 0.75), tol=1e-5):
    for val in values:
        if close_mod1(v, val, tol):
            return True, val
    return False, None


def unique_frac_coords(coords, tol=1e-8):
    unique = []
    for c in coords:
        c = wrap01(c)
        if not any(close_frac(c, u, tol) for u in unique):
            unique.append(c)
    return unique


def expand_site_by_spacegroup(site_frac, sg_symbol, tol=1e-8):
    sg = SpaceGroup(sg_symbol)
    coords = []
    for op in sg.symmetry_ops:
        coords.append(wrap01(op.operate(site_frac)))
    return unique_frac_coords(coords, tol=tol)


def get_lattice_dof_from_crystal_system(crystal_system):
    cs = crystal_system.lower()

    if cs == "triclinic":
        return {
            "free": ["a", "b", "c", "alpha", "beta", "gamma"],
            "constraints": [],
        }

    if cs == "monoclinic":
        return {
            "free": ["a", "b", "c", "beta"],
            "constraints": ["alpha = 90", "gamma = 90"],
        }

    if cs == "orthorhombic":
        return {
            "free": ["a", "b", "c"],
            "constraints": ["alpha = beta = gamma = 90"],
        }

    if cs == "tetragonal":
        return {
            "free": ["a", "c"],
            "constraints": ["a = b", "alpha = beta = gamma = 90"],
        }

    if cs in ("trigonal", "hexagonal"):
        return {
            "free": ["a", "c"],
            "constraints": ["a = b", "alpha = beta = 90", "gamma = 120"],
        }

    if cs == "cubic":
        return {
            "free": ["a"],
            "constraints": ["a = b = c", "alpha = beta = gamma = 90"],
        }

    raise ValueError(f"Unknown crystal system: {crystal_system}")


def get_symmetry_signature(structure, symprec, angle_tolerance):
    sga = SpacegroupAnalyzer(
        structure,
        symprec=symprec,
        angle_tolerance=angle_tolerance,
    )
    symm = sga.get_symmetrized_structure()

    site_info = []
    for i, sites in enumerate(symm.equivalent_sites):
        site0 = sites[0]
        site_info.append({
            "species": site0.species_string,
            "wyckoff": symm.wyckoff_symbols[i],
            "multiplicity": len(sites),
            "frac": wrap01(site0.frac_coords),
        })

    return {
        "sg_symbol": sga.get_space_group_symbol(),
        "sg_number": sga.get_space_group_number(),
        "crystal_system": sga.get_crystal_system(),
        "site_info": site_info,
    }


def same_symmetry_signature(ref, test):
    if ref["sg_number"] != test["sg_number"]:
        return False

    if len(ref["site_info"]) != len(test["site_info"]):
        return False

    ref_sig = sorted(
        (x["species"], x["wyckoff"], x["multiplicity"])
        for x in ref["site_info"]
    )
    test_sig = sorted(
        (x["species"], x["wyckoff"], x["multiplicity"])
        for x in test["site_info"]
    )

    return ref_sig == test_sig


def infer_value_based_constraints(frac, tol=1e-5):
    labels = ["x", "y", "z"]
    frac = wrap01(frac)

    fixed = {}
    for i, lab in enumerate(labels):
        ok, val = is_fixed_value(frac[i], tol=tol)
        if ok:
            fixed[lab] = val

    equal_pairs = []
    for i in range(3):
        for j in range(i + 1, 3):
            if labels[i] not in fixed and labels[j] not in fixed:
                if close_mod1(frac[i], frac[j], tol):
                    equal_pairs.append((labels[i], labels[j]))

    return fixed, equal_pairs


def find_site_indices_for_group(structure, group_sites, tol=1e-5):
    indices = []
    used = set()

    for gs in group_sites:
        target_frac = wrap01(gs.frac_coords)
        found = None

        for i, site in enumerate(structure):
            if i in used:
                continue
            if site.species != gs.species:
                continue
            if close_frac(wrap01(site.frac_coords), target_frac, tol):
                found = i
                break

        if found is None:
            raise RuntimeError(
                f"Could not match equivalent site {gs.species_string} "
                f"at {target_frac} to original structure."
            )

        indices.append(found)
        used.add(found)

    return indices


def match_coords_to_indices(structure, indices, new_coords, tol=1e-4):
    """
    元の等価サイト群 indices と、摂動後に展開した new_coords を対応付ける。
    元座標に最も近いものから貪欲に対応させる。
    """
    remaining = list(new_coords)
    matched = []

    for idx in indices:
        old_frac = wrap01(structure[idx].frac_coords)

        distances = [frac_dist(old_frac, c) for c in remaining]
        j = int(np.argmin(distances))

        matched.append(wrap01(remaining[j]))
        remaining.pop(j)

    return matched


def perturb_equivalent_group_coordinate(
    structure,
    group_indices,
    representative_frac,
    coord_index,
    delta,
    sg_symbol,
    tol=1e-8,
):
    """
    代表サイトを微小変位し、空間群展開で等価サイト群全体を作り直す。
    その座標で元の等価サイト群全体を同時に置き換える。
    """
    rep2 = wrap01(representative_frac.copy())
    rep2[coord_index] += delta
    rep2 = wrap01(rep2)

    expanded = expand_site_by_spacegroup(rep2, sg_symbol, tol=tol)

    if len(expanded) != len(group_indices):
        return None

    new_coords = match_coords_to_indices(
        structure=structure,
        indices=group_indices,
        new_coords=expanded,
        tol=tol,
    )

    s2 = structure.copy()

    for idx, frac in zip(group_indices, new_coords):
        s2.replace(
            idx,
            s2[idx].species,
            frac,
            coords_are_cartesian=False,
        )

    return s2


def analyze_independent_sites(structure, symprec, angle_tolerance, tol, delta):
    ref = get_symmetry_signature(structure, symprec, angle_tolerance)

    sga = SpacegroupAnalyzer(
        structure,
        symprec=symprec,
        angle_tolerance=angle_tolerance,
    )
    symm = sga.get_symmetrized_structure()

    labels = ["x", "y", "z"]
    results = []

    for group_id, sites in enumerate(symm.equivalent_sites):
        site0 = sites[0]
        frac = wrap01(site0.frac_coords)
        wyckoff = symm.wyckoff_symbols[group_id]

        group_indices = find_site_indices_for_group(
            structure,
            sites,
            tol=tol,
        )

        fixed, equal_pairs = infer_value_based_constraints(frac, tol=tol)

        perturb_free = []
        perturb_fixed = []

        for coord_index, lab in enumerate(labels):
            s2 = perturb_equivalent_group_coordinate(
                structure=structure,
                group_indices=group_indices,
                representative_frac=frac,
                coord_index=coord_index,
                delta=delta,
                sg_symbol=ref["sg_symbol"],
                tol=tol,
            )

            if s2 is None:
                ok = False
            else:
                try:
                    test = get_symmetry_signature(
                        s2,
                        symprec=symprec,
                        angle_tolerance=angle_tolerance,
                    )
                    ok = same_symmetry_signature(ref, test)
                except Exception:
                    ok = False

            if ok:
                perturb_free.append(lab)
            else:
                perturb_fixed.append(lab)

        expanded0 = expand_site_by_spacegroup(
            frac,
            ref["sg_symbol"],
            tol=tol,
        )

        results.append({
            "group_id": group_id,
            "site_indices": group_indices,
            "species": site0.species_string,
            "wyckoff": wyckoff,
            "multiplicity": len(sites),
            "frac": frac,
            "fixed_by_value": fixed,
            "equal_by_value": equal_pairs,
            "free_by_perturb": perturb_free,
            "fixed_by_perturb": perturb_fixed,
            "expanded_coords": expanded0,
        })

    return ref, results


def print_coords(coords, indent="        "):
    for i, c in enumerate(coords):
        c = wrap01(c)
        print(
            f"{indent}op[{i:02d}] : "
            f"({c[0]:.8f}, {c[1]:.8f}, {c[2]:.8f})"
        )


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Infer lattice and atomic coordinate degrees of freedom "
            "from CIF using pymatgen."
        )
    )
    parser.add_argument("cif", help="Input CIF file")
    parser.add_argument("--symprec", type=float, default=1e-4)
    parser.add_argument("--angle-tolerance", type=float, default=5.0)
    parser.add_argument("--tol", type=float, default=1e-6)
    parser.add_argument("--delta", type=float, default=1e-3)
    parser.add_argument(
        "--print-expanded",
        type=int,
        default=1,
        choices=[0, 1],
        help="Print expanded equivalent coordinates by space group.",
    )
    args = parser.parse_args()

    if args.delta <= args.symprec:
        print("WARNING: delta should usually be larger than symprec.")
        print(f"         symprec = {args.symprec:g}")
        print(f"         delta   = {args.delta:g}")
        print("         recommended: delta >= 10 * symprec")
        print()

    structure = Structure.from_file(args.cif)

    ref, site_results = analyze_independent_sites(
        structure=structure,
        symprec=args.symprec,
        angle_tolerance=args.angle_tolerance,
        tol=args.tol,
        delta=args.delta,
    )

    lattice_dof = get_lattice_dof_from_crystal_system(ref["crystal_system"])

    print("====================================")
    print("Symmetry")
    print("====================================")
    print(f"Space group    : {ref['sg_symbol']} ({ref['sg_number']})")
    print(f"Crystal system : {ref['crystal_system']}")
    print()

    print("====================================")
    print("Lattice degrees of freedom")
    print("====================================")
    print("Free lattice parameters:")
    print("  " + ", ".join(lattice_dof["free"]))

    if lattice_dof["constraints"]:
        print("Constraints:")
        for c in lattice_dof["constraints"]:
            print(f"  {c}")
    else:
        print("Constraints: none")

    print()

    print("====================================")
    print("Independent atomic sites")
    print("====================================")

    for r in site_results:
        frac = r["frac"]

        print(f"[{r['group_id']}] species      : {r['species']}")
        print(f"    site indices : {r['site_indices']}")
        print(f"    Wyckoff      : {r['wyckoff']}")
        print(f"    multiplicity : {r['multiplicity']}")
        print(
            "    frac coords  : "
            f"({frac[0]:.8f}, {frac[1]:.8f}, {frac[2]:.8f})"
        )

        if r["fixed_by_value"]:
            s = ", ".join(
                f"{k}={v:.8g}" for k, v in r["fixed_by_value"].items()
            )
        else:
            s = "none"
        print(f"    fixed by value : {s}")

        if r["equal_by_value"]:
            s = ", ".join(f"{a}={b}" for a, b in r["equal_by_value"])
        else:
            s = "none"
        print(f"    equal by value : {s}")

        if r["free_by_perturb"]:
            s = ", ".join(r["free_by_perturb"])
        else:
            s = "none"
        print(f"    free by perturbation  : {s}")

        if r["fixed_by_perturb"]:
            s = ", ".join(r["fixed_by_perturb"])
        else:
            s = "none"
        print(f"    fixed by perturbation : {s}")

        print(f"    estimated coordinate dof : {len(r['free_by_perturb'])}")

        if args.print_expanded:
            print("    expanded coords by space group:")
            print_coords(r["expanded_coords"])

        print()


if __name__ == "__main__":
    main()