#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Vibrational irreps from a molecule (XYZ) using point-group character tables.

- Input point group in Schoenflies-like symbols (e.g., C2v, D3d, Td, Oh, C3h, C4h, Ch, I, Ih).
- Uses pymatgen (Hermann–Mauguin) when available; otherwise falls back to an abstract-mode
  that uses known class sizes and a general-position assumption for Γ_3N.
- Builds Γ_3N from fixed-atom counting (site-basis) or general-position assumption, subtracts
  translations/rotations, and decomposes to irreps with character orthogonality.

Added groups:
  D2, D2h, D3, D3d, D3h, D4, D4h, D2d, D6, D6h,
  T, Th, Td, O, Oh,
  C3h, C4h(=4/m), Ch(=m), I(icosahedral), Ih(full icosahedral).

Usage:
  python vib_irreps.py --pg Td  --xyz CH4.xyz
  python vib_irreps.py --pg Oh  --xyz SF6.xyz
  python vib_irreps.py --pg D6h --xyz benzene.xyz
  python vib_irreps.py --pg C4h --xyz something.xyz
  python vib_irreps.py --pg I   --xyz C60.xyz   (abstract mode)
  python vib_irreps.py --pg C3h --xyz molecule.xyz (abstract mode)
"""

import sys
import argparse
import numpy as np
from typing import Dict, List, Tuple, Optional
from collections import Counter, defaultdict

from pymatgen.symmetry.groups import PointGroup

# ===================== Schoenflies -> Hermann–Mauguin =====================
SCH_TO_HM = {
    # basic
    "C1": "1", "CI": "-1", "Ci": "-1", "CS": "m", "Cs": "m", "Ch": "m",
    "C2": "2", "C3": "3", "C4": "4", "C6": "6",
    "C2V": "mm2", "C3V": "3m", "C4V": "4mm", "C6V": "6mm",
    "C2v": "mm2", "C3v": "3m", "C4v": "4mm", "C6v": "6mm",
    "C2h": "2/m",

    # D families
    "D2":  "222",
    "D2H": "mmm", "D2h": "mmm",
    "D3":  "32",
    "D3D": "-3m", "D3d": "-3m",
    "D3H": "-6m2", "D3h": "-6m2",
    "D4":  "422",
    "D4H": "4/mmm", "D4h": "4/mmm",
    "D2D": "-42m",  "D2d": "-42m",
    "D6":  "622",
    "D6H": "6/mmm", "D6h": "6/mmm",

    # cubic
    "T":  "23",
    "TD": "-43m", "Td": "-43m",             # ← 修正済み
    "TH": "m-3",  "Th": "m-3",
    "O":  "432",
    "OH": "m-3m", "Oh": "m-3m",
    "Od": "m-3m",                             # alias

    # C4h is crystallographic: 4/m
    "C4h": "4/m",

    # non-crystallographic, keep Schoenflies (handled by abstract mode)
    "C3h": "C3h",
    "I":   "I",
    "Ih":  "Ih",
}

def to_hm(pg_symbol: str) -> str:
    s = pg_symbol.strip()
    return SCH_TO_HM.get(s, s)

# ===================== Character tables =====================
# 注: 一部の群（C4h）は複素 1 次元既約を用います（群が可換：8 個の 1D）。
#     分解内積では複素共役を使うので問題ありません。

def gen_C4h_table():
    """
    C4h ≡ 4/m. Abelian of order 8 with generators r=C4, u=i (inversion), σh = u*r^2.
    Classes = elements: E, C4, C4^3, C2, i, S4, S4^3, σh.
    Irreps: χ_{a,par}, a∈{0,1,2,3}, par∈{+1(g), -1(u)} with
            χ(r^k) = exp(i*pi/2 * a*k), χ(u*r^k) = par * exp(i*pi/2 * a*k).
    Labels: Γ0g, Γ1g, Γ2g, Γ3g, Γ0u, Γ1u, Γ2u, Γ3u
    """
    classes = ["E","C4","C4^3","C2","i","S4","S4^3","σh"]
    irreps = {}
    iunit = 1j
    for a in (0,1,2,3):
        for par, tag in ((+1,"g"),(-1,"u")):
            lab = f"Γ{a}{tag}"
            def phase(k): return np.exp(1j*np.pi/2 * a * k)
            irreps[lab] = {
                "E":   1.0,
                "C4":  phase(1),
                "C4^3":phase(3),
                "C2":  phase(2),
                "i":   par * 1.0,
                "S4":  par * phase(1),
                "S4^3":par * phase(3),
                "σh":  par * phase(2),
            }
    return {"classes": classes, "irreps": irreps}

# 黄金比（I, Ih 用）
phi  = (1 + 5**0.5)/2
phip = (1 - 5**0.5)/2

CHAR_TABLES: Dict[str, Dict[str, Dict[str, complex]]] = {
    # ---- simple C, Cnv (既存) ----
    "C1": {"classes": ["E"], "irreps": {"A": {"E": 1}}},
    "Ci": {"classes": ["E","i"], "irreps": {"Ag":{"E":1,"i":1}, "Au":{"E":1,"i":-1}}},
    "Cs": {"classes": ["E","σ"], "irreps": {"A'":{"E":1,"σ":1}, "A''":{"E":1,"σ":-1}}},
    "Ch": {"classes": ["E","σ"], "irreps": {"A'":{"E":1,"σ":1}, "A''":{"E":1,"σ":-1}}},  # alias
    "C2": {"classes":["E","C2"], "irreps":{"A":{"E":1,"C2":1}, "B":{"E":1,"C2":-1}}},
    "C3": {"classes":["E","C3","C3^2"], "irreps":{
        "A":{"E":1,"C3":1,"C3^2":1}, "E":{"E":2,"C3":-1,"C3^2":-1}
    }},
    "C4": {"classes":["E","C4","C2","C4^3"], "irreps":{
        "A":{"E":1,"C4":1,"C2":1,"C4^3":1},
        "B":{"E":1,"C4":-1,"C2":1,"C4^3":-1},
        "E":{"E":2,"C4":0,"C2":-2,"C4^3":0}
    }},
    "C6": {"classes":["E","C6","C3","C2","C3^2","C6^5"], "irreps":{
        "A":{"E":1,"C6":1,"C3":1,"C2":1,"C3^2":1,"C6^5":1},
        "B":{"E":1,"C6":-1,"C3":1,"C2":-1,"C3^2":1,"C6^5":-1},
        "E1":{"E":2,"C6":1,"C3":-1,"C2":-2,"C3^2":-1,"C6^5":1},
        "E2":{"E":2,"C6":-1,"C3":-1,"C2":2,"C3^2":-1,"C6^5":-1}
    }},
    "C2v": {"classes":["E","C2","σv","σv'"], "irreps":{
        "A1":{"E":1,"C2":1,"σv":1,"σv'":1},
        "A2":{"E":1,"C2":1,"σv":-1,"σv'":-1},
        "B1":{"E":1,"C2":-1,"σv":1,"σv'":-1},
        "B2":{"E":1,"C2":-1,"σv":-1,"σv'":1}
    }},
    "C3v": {"classes":["E","C3","C3^2","σv"], "irreps":{
        "A1":{"E":1,"C3":1,"C3^2":1,"σv":1},
        "A2":{"E":1,"C3":1,"C3^2":1,"σv":-1},
        "E" :{"E":2,"C3":-1,"C3^2":-1,"σv":0}
    }},
    "C4v": {"classes":["E","C4","C4^3","C2","σv","σd"], "irreps":{
        "A1":{"E":1,"C4":1,"C4^3":1,"C2":1,"σv":1,"σd":1},
        "A2":{"E":1,"C4":1,"C4^3":1,"C2":1,"σv":-1,"σd":-1},
        "B1":{"E":1,"C4":-1,"C4^3":-1,"C2":1,"σv":1,"σd":-1},
        "B2":{"E":1,"C4":-1,"C4^3":-1,"C2":1,"σv":-1,"σd":1},
        "E" :{"E":2,"C4":0,"C4^3":0,"C2":-2,"σv":0,"σd":0}
    }},
    "C6v": {"classes":["E","C6","C6^5","C3","C3^2","C2","σv","σd"], "irreps":{
        "A1":{"E":1,"C6":1,"C6^5":1,"C3":1,"C3^2":1,"C2":1,"σv":1,"σd":1},
        "A2":{"E":1,"C6":1,"C6^5":1,"C3":1,"C3^2":1,"C2":1,"σv":-1,"σd":-1},
        "B1":{"E":1,"C6":-1,"C6^5":-1,"C3":1,"C3^2":1,"C2":-1,"σv":1,"σd":-1},
        "B2":{"E":1,"C6":-1,"C6^5":-1,"C3":1,"C3^2":1,"C2":-1,"σv":-1,"σd":1},
        "E1":{"E":2,"C6":1,"C6^5":1,"C3":-1,"C3^2":-1,"C2":-2,"σv":0,"σd":0},
        "E2":{"E":2,"C6":-1,"C6^5":-1,"C3":-1,"C3^2":-1,"C2":2,"σv":0,"σd":0}
    }},

    # ---- C2h ----
    "C2h":{"classes":["E","C2(z)","i","σh"], "irreps":{
        "Ag":{"E":1,"C2(z)":1,"i":1,"σh":1},
        "Bg":{"E":1,"C2(z)":-1,"i":1,"σh":-1},
        "Au":{"E":1,"C2(z)":1,"i":-1,"σh":-1},
        "Bu":{"E":1,"C2(z)":-1,"i":-1,"σh":1},
    }},

    # ---- D / trigonal / tetragonal / hexagonal ----
    "D2":{"classes":["E","C2(x)","C2(y)","C2(z)"], "irreps":{
        "A" :{"E":1,"C2(x)":1,"C2(y)":1,"C2(z)":1},
        "B1":{"E":1,"C2(x)":1,"C2(y)":-1,"C2(z)":-1},
        "B2":{"E":1,"C2(x)":-1,"C2(y)":1,"C2(z)":-1},
        "B3":{"E":1,"C2(x)":-1,"C2(y)":-1,"C2(z)":1},
    }},
    "D2h":{"classes":["E","C2(x)","C2(y)","C2(z)","i","σ(xy)","σ(xz)","σ(yz)"], "irreps":{
        "Ag":{"E":1,"C2(x)":1,"C2(y)":1,"C2(z)":1,"i":1,"σ(xy)":1,"σ(xz)":1,"σ(yz)":1},
        "B1g":{"E":1,"C2(x)":1,"C2(y)":-1,"C2(z)":-1,"i":1,"σ(xy)":1,"σ(xz)":-1,"σ(yz)":-1},
        "B2g":{"E":1,"C2(x)":-1,"C2(y)":1,"C2(z)":-1,"i":1,"σ(xy)":-1,"σ(xz)":1,"σ(yz)":-1},
        "B3g":{"E":1,"C2(x)":-1,"C2(y)":-1,"C2(z)":1,"i":1,"σ(xy)":-1,"σ(xz)":-1,"σ(yz)":1},
        "Au":{"E":1,"C2(x)":1,"C2(y)":1,"C2(z)":1,"i":-1,"σ(xy)":-1,"σ(xz)":-1,"σ(yz)":-1},
        "B1u":{"E":1,"C2(x)":1,"C2(y)":-1,"C2(z)":-1,"i":-1,"σ(xy)":-1,"σ(xz)":1,"σ(yz)":1},
        "B2u":{"E":1,"C2(x)":-1,"C2(y)":1,"C2(z)":-1,"i":-1,"σ(xy)":1,"σ(xz)":-1,"σ(yz)":1},
        "B3u":{"E":1,"C2(x)":-1,"C2(y)":-1,"C2(z)":1,"i":-1,"σ(xy)":1,"σ(xz)":1,"σ(yz)":-1},
    }},
    "D3":{"classes":["E","C3","C2'"], "irreps":{
        "A1":{"E":1,"C3":1,"C2'":1},
        "A2":{"E":1,"C3":1,"C2'":-1},
        "E" :{"E":2,"C3":-1,"C2'":0},
    }},
    "D3d":{"classes":["E","C3","C2'","i","S6","σd"], "irreps":{
        "A1g":{"E":1,"C3":1,"C2'":1,"i":1,"S6":1,"σd":1},
        "A2g":{"E":1,"C3":1,"C2'":-1,"i":1,"S6":-1,"σd":-1},
        "Eg" :{"E":2,"C3":-1,"C2'":0,"i":2,"S6":0,"σd":0},
        "A1u":{"E":1,"C3":1,"C2'":1,"i":-1,"S6":-1,"σd":-1},
        "A2u":{"E":1,"C3":1,"C2'":-1,"i":-1,"S6":1,"σd":1},
        "Eu" :{"E":2,"C3":-1,"C2'":0,"i":-2,"S6":0,"σd":0},
    }},
    "D3h":{"classes":["E","C3","C2'","σh","S3","σv"], "irreps":{
        "A1'": {"E":1,"C3":1,"C2'":1,"σh":1,"S3":1,"σv":1},
        "A2'": {"E":1,"C3":1,"C2'":-1,"σh":1,"S3":1,"σv":-1},
        "E'":  {"E":2,"C3":-1,"C2'":0,"σh":2,"S3":-1,"σv":0},
        "A1''":{"E":1,"C3":1,"C2'":1,"σh":-1,"S3":-1,"σv":-1},
        "A2''":{"E":1,"C3":1,"C2'":-1,"σh":-1,"S3":-1,"σv":1},
        "E''": {"E":2,"C3":-1,"C2'":0,"σh":-2,"S3":1,"σv":0},
    }},
    "D4":{"classes":["E","C4","C2(z)","C2'","C2''"], "irreps":{
        "A1":{"E":1,"C4":1,"C2(z)":1,"C2'":1,"C2''":1},
        "A2":{"E":1,"C4":1,"C2(z)":1,"C2'":-1,"C2''":-1},
        "B1":{"E":1,"C4":-1,"C2(z)":1,"C2'":1,"C2''":-1},
        "B2":{"E":1,"C4":-1,"C2(z)":1,"C2'":-1,"C2''":1},
        "E" :{"E":2,"C4":0,"C2(z)":-2,"C2'":0,"C2''":0},
    }},
    "D2d":{"classes":["E","S4","C2(z)","C2'","σd"], "irreps":{
        "A1":{"E":1,"S4":1,"C2(z)":1,"C2'":1,"σd":1},
        "A2":{"E":1,"S4":1,"C2(z)":1,"C2'":-1,"σd":-1},
        "B1":{"E":1,"S4":-1,"C2(z)":1,"C2'":1,"σd":-1},
        "B2":{"E":1,"S4":-1,"C2(z)":1,"C2'":-1,"σd":1},
        "E" :{"E":2,"S4":0,"C2(z)":-2,"C2'":0,"σd":0},
    }},
    "D4h":{"classes":["E","C4","C2(z)","C2'","C2''","i","S4","σh","σv","σd"], "irreps":{
        "A1g":{"E":1,"C4":1,"C2(z)":1,"C2'":1,"C2''":1,"i":1,"S4":1,"σh":1,"σv":1,"σd":1},
        "A2g":{"E":1,"C4":1,"C2(z)":1,"C2'":-1,"C2''":-1,"i":1,"S4":1,"σh":1,"σv":-1,"σd":-1},
        "B1g":{"E":1,"C4":-1,"C2(z)":1,"C2'":1,"C2''":-1,"i":1,"S4":-1,"σh":1,"σv":1,"σd":-1},
        "B2g":{"E":1,"C4":-1,"C2(z)":1,"C2'":-1,"C2''":1,"i":1,"S4":-1,"σh":1,"σv":-1,"σd":1},
        "Eg" :{"E":2,"C4":0,"C2(z)":-2,"C2'":0,"C2''":0,"i":2,"S4":0,"σh":2,"σv":0,"σd":0},
        "A1u":{"E":1,"C4":1,"C2(z)":1,"C2'":1,"C2''":1,"i":-1,"S4":-1,"σh":-1,"σv":-1,"σd":-1},
        "A2u":{"E":1,"C4":1,"C2(z)":1,"C2'":-1,"C2''":-1,"i":-1,"S4":-1,"σh":-1,"σv":1,"σd":1},
        "B1u":{"E":1,"C4":-1,"C2(z)":1,"C2'":1,"C2''":-1,"i":-1,"S4":1,"σh":-1,"σv":-1,"σd":1},
        "B2u":{"E":1,"C4":-1,"C2(z)":1,"C2'":-1,"C2''":1,"i":-1,"S4":1,"σh":-1,"σv":1,"σd":-1},
        "Eu" :{"E":2,"C4":0,"C2(z)":-2,"C2'":0,"C2''":0,"i":-2,"S4":0,"σh":-2,"σv":0,"σd":0},
    }},
    "D6":{"classes":["E","C6","C3","C2(z)","C2'","C2''"], "irreps":{
        "A1":{"E":1,"C6":1,"C3":1,"C2(z)":1,"C2'":1,"C2''":1},
        "A2":{"E":1,"C6":1,"C3":1,"C2(z)":1,"C2'":-1,"C2''":-1},
        "B1":{"E":1,"C6":-1,"C3":1,"C2(z)":-1,"C2'":1,"C2''":-1},
        "B2":{"E":1,"C6":-1,"C3":1,"C2(z)":-1,"C2'":-1,"C2''":1},
        "E1":{"E":2,"C6":1,"C3":-1,"C2(z)":-2,"C2'":0,"C2''":0},
        "E2":{"E":2,"C6":-1,"C3":-1,"C2(z)":2,"C2'":0,"C2''":0},
    }},
    "D6h":{"classes":["E","C6","C3","C2(z)","C2'","C2''","i","S3","S6","σh","σd","σv"], "irreps":{
        "A1g":{"E":1,"C6":1,"C3":1,"C2(z)":1,"C2'":1,"C2''":1,"i":1,"S3":1,"S6":1,"σh":1,"σd":1,"σv":1},
        "A2g":{"E":1,"C6":1,"C3":1,"C2(z)":1,"C2'":-1,"C2''":-1,"i":1,"S3":1,"S6":1,"σh":1,"σd":-1,"σv":-1},
        "B1g":{"E":1,"C6":-1,"C3":1,"C2(z)":-1,"C2'":1,"C2''":-1,"i":1,"S3":-1,"S6":1,"σh":1,"σd":1,"σv":-1},
        "B2g":{"E":1,"C6":-1,"C3":1,"C2(z)":-1,"C2'":-1,"C2''":1,"i":1,"S3":-1,"S6":1,"σh":1,"σd":-1,"σv":1},
        "E1g":{"E":2,"C6":1,"C3":-1,"C2(z)":-2,"C2'":0,"C2''":0,"i":2,"S3":1,"S6":-1,"σh":2,"σd":0,"σv":0},
        "E2g":{"E":2,"C6":-1,"C3":-1,"C2(z)":2,"C2'":0,"C2''":0,"i":2,"S3":-1,"S6":1,"σh":2,"σd":0,"σv":0},
        "A1u":{"E":1,"C6":1,"C3":1,"C2(z)":1,"C2'":1,"C2''":1,"i":-1,"S3":-1,"S6":-1,"σh":-1,"σd":-1,"σv":-1},
        "A2u":{"E":1,"C6":1,"C3":1,"C2(z)":1,"C2'":-1,"C2''":-1,"i":-1,"S3":-1,"S6":-1,"σh":-1,"σd":1,"σv":1},
        "B1u":{"E":1,"C6":-1,"C3":1,"C2(z)":-1,"C2'":1,"C2''":-1,"i":-1,"S3":1,"S6":-1,"σh":-1,"σd":-1,"σv":1},
        "B2u":{"E":1,"C6":-1,"C3":1,"C2(z)":-1,"C2'":-1,"C2''":1,"i":-1,"S3":1,"S6":-1,"σh":-1,"σd":1,"σv":-1},
        "E1u":{"E":2,"C6":1,"C3":-1,"C2(z)":-2,"C2'":0,"C2''":0,"i":-2,"S3":-1,"S6":1,"σh":-2,"σd":0,"σv":0},
        "E2u":{"E":2,"C6":-1,"C3":-1,"C2(z)":2,"C2'":0,"C2''":0,"i":-2,"S3":1,"S6":-1,"σh":-2,"σd":0,"σv":0},
    }},

    # ---- cubic ----
    "T":{"classes":["E","C3","C2"], "irreps":{
        "A":{"E":1,"C3":1,"C2":1},
        "E":{"E":2,"C3":-1,"C2":2},
        "T":{"E":3,"C3":0,"C2":-1},
    }},
    "Td":{"classes":["E","C3","C2","S4","σd"], "irreps":{
        "A1":{"E":1,"C3":1,"C2":1,"S4":1,"σd":1},
        "A2":{"E":1,"C3":1,"C2":1,"S4":-1,"σd":-1},
        "E" :{"E":2,"C3":-1,"C2":2,"S4":0,"σd":0},
        "T1":{"E":3,"C3":0,"C2":-1,"S4":1,"σd":-1},
        "T2":{"E":3,"C3":0,"C2":-1,"S4":-1,"σd":1},
    }},
    "Th":{"classes":["E","C3","C2","i","S6"], "irreps":{
        "Ag":{"E":1,"C3":1,"C2":1,"i":1,"S6":1},
        "Au":{"E":1,"C3":1,"C2":1,"i":-1,"S6":-1},
        "Eg":{"E":2,"C3":-1,"C2":2,"i":2,"S6":0},
        "Eu":{"E":2,"C3":-1,"C2":2,"i":-2,"S6":0},
        "Tg":{"E":3,"C3":0,"C2":-1,"i":3,"S6":-1},
        "Tu":{"E":3,"C3":0,"C2":-1,"i":-3,"S6":1},
    }},
    "O":{"classes":["E","C3","C4","C2(ax)","C2(di)"], "irreps":{
        "A1":{"E":1,"C3":1,"C4":1,"C2(ax)":1,"C2(di)":1},
        "A2":{"E":1,"C3":1,"C4":-1,"C2(ax)":1,"C2(di)":-1},
        "E" :{"E":2,"C3":-1,"C4":0,"C2(ax)":2,"C2(di)":0},
        "T1":{"E":3,"C3":0,"C4":1,"C2(ax)":-1,"C2(di)":-1},
        "T2":{"E":3,"C3":0,"C4":-1,"C2(ax)":-1,"C2(di)":1},
    }},
    "Oh":{"classes":["E","C3","C4","C2(ax)","C2(di)","i","S6","S4","σh","σd"], "irreps":{
        "A1g":{"E":1,"C3":1,"C4":1,"C2(ax)":1,"C2(di)":1,"i":1,"S6":1,"S4":1,"σh":1,"σd":1},
        "A2g":{"E":1,"C3":1,"C4":-1,"C2(ax)":1,"C2(di)":-1,"i":1,"S6":-1,"S4":-1,"σh":1,"σd":-1},
        "Eg" :{"E":2,"C3":-1,"C4":0,"C2(ax)":2,"C2(di)":0,"i":2,"S6":0,"S4":0,"σh":2,"σd":0},
        "T1g":{"E":3,"C3":0,"C4":1,"C2(ax)":-1,"C2(di)":-1,"i":3,"S6":-1,"S4":1,"σh":-1,"σd":-1},
        "T2g":{"E":3,"C3":0,"C4":-1,"C2(ax)":-1,"C2(di)":1,"i":3,"S6":1,"S4":-1,"σh":-1,"σd":1},
        "A1u":{"E":1,"C3":1,"C4":1,"C2(ax)":1,"C2(di)":1,"i":-1,"S6":-1,"S4":-1,"σh":-1,"σd":-1},
        "A2u":{"E":1,"C3":1,"C4":-1,"C2(ax)":1,"C2(di)":-1,"i":-1,"S6":1,"S4":1,"σh":-1,"σd":1},
        "Eu" :{"E":2,"C3":-1,"C4":0,"C2(ax)":2,"C2(di)":0,"i":-2,"S6":0,"S4":0,"σh":-2,"σd":0},
        "T1u":{"E":3,"C3":0,"C4":1,"C2(ax)":-1,"C2(di)":-1,"i":-3,"S6":1,"S4":-1,"σh":1,"σd":1},
        "T2u":{"E":3,"C3":0,"C4":-1,"C2(ax)":-1,"C2(di)":1,"i":-3,"S6":-1,"S4":1,"σh":1,"σd":-1},
    }},

    # ---- C3h (real-valued 4-class table; 2D E', E'') ----
    "C3h": {"classes": ["E","C3","σh","S3"], "irreps": {
        "A'":  {"E":1, "C3": 1, "σh": 1, "S3": 1},
        "A''": {"E":1, "C3": 1, "σh":-1, "S3":-1},
        "E'":  {"E":2, "C3":-1, "σh": 2, "S3":-1},
        "E''": {"E":2, "C3":-1, "σh":-2, "S3": 1},
    }},

    # ---- I (icosahedral, order 60) ----
    # Classes: E (1), 12C5, 12C5^2, 20C3, 15C2
    "I": {"classes": ["E","C5","C5^2","C3","C2"], "irreps": {
        "A":  {"E":1, "C5":1,    "C5^2":1,    "C3":1,  "C2":1},
        "T1": {"E":3, "C5":phi,  "C5^2":phip, "C3":0,  "C2":-1},
        "T2": {"E":3, "C5":phip, "C5^2":phi,  "C3":0,  "C2":-1},
        "G":  {"E":4, "C5":-1,   "C5^2":-1,   "C3":1,  "C2":0},
        "H":  {"E":5, "C5":0,    "C5^2":0,    "C3":-1, "C2":1},
    }},

    # ---- Ih (full icosahedral, order 120) ----
    # Classes: add inversion counterparts: i, 12S10, 12S10^3, 20S6, 15S2
    # χ_g on improper = + χ_I(of corresponding rotation), χ_u = - χ_I(...)
    "Ih": {"classes": ["E","C5","C5^2","C3","C2","i","S10","S10^3","S6","S2"], "irreps": {}},
}

# Fill Ih irreps from I
def _fill_Ih():
    base = CHAR_TABLES["I"]["irreps"]
    irrIh = {}
    for name, row in base.items():
        # gerade
        gname = name + "g"
        irrIh[gname] = dict(row)
        irrIh[gname].update({
            "i":  +row["E"],
            "S10": +row["C5"],
            "S10^3": +row["C5^2"],
            "S6": +row["C3"],
            "S2": +row["C2"],
        })
        # ungerade
        uname = name + "u"
        irrIh[uname] = dict(row)
        irrIh[uname].update({
            "i":  -row["E"],
            "S10": -row["C5"],
            "S10^3": -row["C5^2"],
            "S6": -row["C3"],
            "S2": -row["C2"],
        })
    CHAR_TABLES["Ih"]["irreps"] = irrIh

_fill_Ih()
CHAR_TABLES["C4h"] = gen_C4h_table()

SUPPORTED = set(CHAR_TABLES.keys())

# ===== Abstract-mode class sizes (when pymatgen lacks the group) =====
ABSTRACT_CLASS_SIZES = {
    # C3h: order 6 with merged real classes (E, 2C3, σh, 2S3)
    "C3h": {"E":1, "C3":2, "σh":1, "S3":2},

    # I: A5, order 60
    "I":   {"E":1, "C5":12, "C5^2":12, "C3":20, "C2":15},

    # Ih: order 120
    "Ih":  {"E":1, "C5":12, "C5^2":12, "C3":20, "C2":15,
            "i":1, "S10":12, "S10^3":12, "S6":20, "S2":15},
}

# ===================== Numerics helper =====================
def orth(R: np.ndarray) -> np.ndarray:
    U, _, Vt = np.linalg.svd(R)
    R_ = U @ Vt
    if np.linalg.det(R_) < 0:
        U[:, -1] *= -1
        R_ = U @ Vt
    return R_

# ===================== Group ops via pymatgen (with robust fallback) =====================
def pg_ops(pg_symbol: str) -> List[np.ndarray]:
    hm = to_hm(pg_symbol)
    candidates = [hm, hm.replace(" ", ""), hm.replace("−","-"), pg_symbol]
    last_err = None
    for cand in candidates:
        try:
            ops = PointGroup(cand).symmetry_ops
            return [orth(op.rotation_matrix) for op in ops]
        except Exception as e:
            last_err = e
            continue
    raise RuntimeError(f"PointGroup lookup failed for '{pg_symbol}' (tried {candidates}): {last_err}")

def pg_order(pg_symbol: str) -> int:
    try:
        return len(pg_ops(pg_symbol))
    except Exception:
        # abstract groups
        if pg_symbol in ABSTRACT_CLASS_SIZES:
            return sum(ABSTRACT_CLASS_SIZES[pg_symbol].values())
        # fallback to table size if provided
        if pg_symbol in SUPPORTED:
            # approximate by assuming unique ops per class label (not used normally)
            return 0
        raise

# ===================== Classification =====================
Z = np.array([0.0, 0.0, 1.0])
X = np.array([1.0, 0.0, 0.0])
Y = np.array([0.0, 1.0, 0.0])
XY_DIAGS = [np.array([1,1,0.0]), np.array([1,-1,0.0]),
            np.array([-1,1,0.0]), np.array([-1,-1,0.0])]
C2_DIAGONALS_3D = [
    np.array([1,1,0.0]), np.array([1,-1,0.0]), np.array([-1,1,0.0]), np.array([-1,-1,0.0]),
    np.array([1,0,1.0]), np.array([1,0,-1.0]), np.array([-1,0,1.0]), np.array([-1,0,-1.0]),
    np.array([0,1,1.0]), np.array([0,1,-1.0]), np.array([0,-1,1.0]), np.array([0,-1,-1.0]),
]

def unit(v):
    v = np.asarray(v, float)
    n = np.linalg.norm(v)
    return v / (n + 1e-15)

def close_to(v, cand_list, tol=0.1):
    v = unit(v)
    best_dot = -1
    for c in cand_list:
        c = unit(c)
        d = abs(np.dot(v, c))
        if d > best_dot:
            best_dot = d
    return best_dot >= (1 - tol)

def rotation_axis(R: np.ndarray, tol=1e-6):
    vals, vecs = np.linalg.eig(R)
    idx = np.argmax(np.real(vals))
    if np.isclose(vals[idx].real, 1.0, atol=1e-5):
        ax = np.real(vecs[:, idx])
        return unit(ax)
    return None

def rotation_angle_from_trace(R: np.ndarray):
    cos_th = float(np.clip((np.trace(R) - 1) / 2, -1, 1))
    th = np.arccos(cos_th)
    return th

def classify_op_for_table(pg: str, R: np.ndarray, tol=1e-6) -> str:
    R = orth(R)
    det = float(np.linalg.det(R))

    # Identity / inversion
    if np.allclose(R, np.eye(3), atol=tol): return "E"
    if np.allclose(R, -np.eye(3), atol=tol): return "i"

    # Mirrors (det = -1 and order 2)
    if det < 0 and np.allclose(R @ R, np.eye(3), atol=1e-6):
        vals, vecs = np.linalg.eig(R)
        idx = np.argmin(np.real(vals))  # eigenvalue ≈ -1 (normal)
        nrm = unit(np.real(vecs[:, idx]))
        nz = abs(nrm[2])
        if nz > 1 - 1e-3:
            return "σh"
        vscore = max(abs(np.dot(nrm, unit(v))) for v in [X,Y,-X,-Y])
        dscore = max(abs(np.dot(nrm, unit(v))) for v in XY_DIAGS)
        if vscore >= dscore:
            return "σv" if abs(np.dot(nrm, X)) >= abs(np.dot(nrm, Y)) else "σv'"
        else:
            return "σd"

    # Proper rotations (det = +1)
    if det > 0:
        th = rotation_angle_from_trace(R)
        if th < 1e-6: return "E"
        n = int(round(2 * np.pi / th))
        ax = rotation_axis(R)

        if n == 2:
            if pg in ("D2","D2h"):
                if close_to(ax, [X,-X]): return "C2(x)"
                if close_to(ax, [Y,-Y]): return "C2(y)"
                if close_to(ax, [Z,-Z]): return "C2(z)"
            if pg in ("C2h",):
                if close_to(ax, [Z,-Z]): return "C2(z)"
            if pg in ("D4","D4h","D2d","C4h","C4v","Oh","O","D6","D6h"):
                if close_to(ax, [Z,-Z]): return "C2(z)"
                if close_to(ax, [X,-X,Y,-Y]): return "C2'"
                if close_to(ax, XY_DIAGS): return "C2''"
                if pg in ("O","Oh"):
                    if close_to(ax, [X,-X,Y,-Y,Z,-Z]): return "C2(ax)"
                    if close_to(ax, C2_DIAGONALS_3D):   return "C2(di)"
            return "C2"

        if n == 3:
            # one class for many groups; use C3 and map C3^2 → C3 if needed
            return "C3" if th <= np.pi/1.5 else "C3^2"

        if n == 4:
            # keep C4 vs C4^3 distinction where needed (e.g. C4h abelian)
            return "C4" if R[0,1] < 0 else "C4^3"

        if n == 5:
            # split C5 vs C5^2 by angle 72° vs 144°
            if np.isclose(th, 2*np.pi/5, atol=1e-3):
                return "C5"
            elif np.isclose(th, 4*np.pi/5, atol=1e-3):
                return "C5^2"
            else:
                return "C5"

        if n == 6:
            if np.isclose(th, np.pi/3, atol=1e-3):   return "C6"
            if np.isclose(th, 2*np.pi/3, atol=1e-3): return "C3^2"
            if np.isclose(th, np.pi, atol=1e-3):     return "C2"
            if np.isclose(th, 4*np.pi/3, atol=1e-3): return "C3"
            if np.isclose(th, 5*np.pi/3, atol=1e-3): return "C6"
        return f"C{n}"

    # Improper rotations: S_n (det = -1, not a mirror)
    if det < 0:
        A = -R
        th = rotation_angle_from_trace(A)  # angle of underlying rotation
        n = int(round(2 * np.pi / th))
        if n == 4:
            # Distinguish S4 vs S4^3 using 2×2 block sign
            if A[0,1] < 0: return "S4"
            else:          return "S4^3"
        if n == 6:
            return "S6"  # common in Th/Oh/D6h
        if n == 3:
            return "S3"  # used in D3h/C3h
        if n == 10:
            # split S10 vs S10^3 by angle ~72° vs ~144° (principal angles)
            if np.isclose(th, 2*np.pi/5, atol=5e-3):      return "S10"
            elif np.isclose(th, 4*np.pi/5, atol=5e-3):    return "S10^3"
            else:                                         return "S10"
        return f"S{n}"

    return "?"

# ===================== XYZ loader =====================
def load_coords_from_xyz(path: str) -> np.ndarray:
    coords = []
    with open(path, "r", encoding="utf-8") as f:
        lines = [ln.strip() for ln in f if ln.strip()]
    start = 0
    try:
        int(lines[0].split()[0]); start = 2
    except Exception:
        start = 0
    for ln in lines[start:]:
        toks = ln.replace(",", " ").split()
        if len(toks) < 4:
            continue
        x, y, z = map(float, toks[-3:])
        coords.append([x, y, z])
    return np.array(coords, dtype=float)

# ===================== Character helpers for Γ_T and Γ_R =====================
def _parse_power(label: str, default_n=None):
    # parse like "C4^3" -> ("C",4,3); "C4" -> ("C",4,1)
    if "^" in label:
        base, p = label.split("^", 1)
        p = int(p)
    else:
        base, p = label, 1
    kind = base[0]  # 'C' or 'S'
    try:
        n = int(base[1:])
    except ValueError:
        n = default_n
    return kind, n, p

def polar_trace_from_label(lab: str) -> float:
    """
    Trace of the 3D polar vector rep (translations) for a single operation class label.
    For rotations C_n^k: tr = 1 + 2 cos(2πk/n)
    For improper S_n^k: tr = 2 cos(2πk/n) - 1
    Mirrors σ*: tr = 1, inversion i: tr = -3, identity E: 3.
    """
    if lab == "E": return 3.0
    if lab == "i": return -3.0
    if lab.startswith("σ"): return 1.0
    if lab.startswith("C"):
        kind, n, p = _parse_power(lab)
        if n is None: return 0.0
        ang = 2*np.pi*(p % n)/n
        return 1.0 + 2.0*np.cos(ang)
    if lab.startswith("S"):
        kind, n, p = _parse_power(lab)
        if n is None: return 0.0
        ang = 2*np.pi*(p % n)/n
        return 2.0*np.cos(ang) - 1.0
    return 0.0

def det_from_label(lab: str) -> int:
    if lab in ("E",): return 1
    if lab.startswith("C"): return 1
    # i, σ*, S* are improper (det = -1)
    return -1

def gamma_3N_characters(coords: np.ndarray, op_mats: List[np.ndarray],
                        labels: List[str], tol=1e-5) -> Dict[str, float]:
    chars = defaultdict(float)
    for R, lab in zip(op_mats, labels):
        chi = 0.0
        for r in coords:
            rp = R @ r
            if np.linalg.norm(rp - r) < tol:
                chi += np.trace(R)
        chars[lab] += float(chi)
    return dict(chars)

def gamma_trans_characters(labels: List[str], class_map: Dict[str, str]) -> Dict[str, float]:
    chars = defaultdict(float)
    for lab in labels:
        key = class_map.get(lab, lab)
        chars[key] += polar_trace_from_label(lab)
    return dict(chars)

def gamma_rot_characters(labels: List[str], class_map: Dict[str, str]) -> Dict[str, float]:
    chars = defaultdict(float)
    for lab in labels:
        key = class_map.get(lab, lab)
        tr = polar_trace_from_label(lab)
        det = det_from_label(lab)
        chars[key] += det * tr
    return dict(chars)

# ===================== Class aggregation & decomposition =====================
def class_aggregation(pg: str, raw_labels: List[str]) -> Tuple[List[str], Dict[str, str]]:
    """
    Map raw op labels to the canonical class labels present in CHAR_TABLES[pg]["classes"].
    """
    classes = CHAR_TABLES[pg]["classes"]
    class_set = set(classes)
    class_map = {}
    for lab in raw_labels:
        if lab in class_set:
            class_map[lab] = lab
            continue
        # Common normalizations:
        if lab == "C2" and "C2(z)" in class_set: class_map[lab] = "C2(z)"; continue
        if lab.startswith("σv") and "σv" in class_set: class_map[lab] = "σv"; continue
        if lab.startswith("σd") and "σd" in class_set: class_map[lab] = "σd"; continue
        if lab in ("C2'", "C2''") and lab in class_set: class_map[lab] = lab; continue
        if lab == "C2(z)" and "C2(z)" in class_set: class_map[lab] = "C2(z)"; continue
        if lab == "C2" and "C2'" in class_set: class_map[lab] = "C2'"; continue
        if lab in ("C2(ax)","C2(di)") and lab in class_set: class_map[lab] = lab; continue
        if lab.startswith("C4") and "C4" in class_set and "C4^3" not in class_set: class_map[lab] = "C4"; continue
        if lab.startswith("C3") and "C3" in class_set: class_map[lab] = "C3"; continue
        if lab.startswith("S3") and "S6" in class_set and "S3" not in class_set: class_map[lab] = "S6"; continue
        if lab in ("S10","S10^3") and lab in class_set: class_map[lab] = lab; continue
        if lab.startswith("S4") and "S4" in class_set and "S4^3" not in class_set: class_map[lab] = "S4"; continue
        if lab.startswith("σ") and "σ" in class_set: class_map[lab] = "σ"; continue
        class_map[lab] = lab
    return classes, class_map

def reduce_to_classes(vec_by_label: Dict[str, float], classes: List[str], class_map: Dict[str, str]) -> Dict[str, float]:
    agg = defaultdict(float)
    for lab, val in vec_by_label.items():
        key = class_map.get(lab, lab)
        agg[key] += val
    return {c: agg.get(c, 0.0) for c in classes}

def decompose(chars: Dict[str, float], pg: str,
             class_sizes_override: Optional[Dict[str,int]] = None,
             order_override: Optional[int] = None) -> Dict[str, int]:
    """
    Decompose class-characters 'chars' into irreps for group 'pg'.
    Uses complex-conjugated characters for orthogonality.
    """
    irreps = CHAR_TABLES[pg]["irreps"]

    if class_sizes_override is not None:
        class_sizes = Counter(class_sizes_override)
        h = order_override if order_override is not None else sum(class_sizes.values())
    else:
        h = pg_order(pg)
        ops = pg_ops(pg)
        raw_labels = [classify_op_for_table(pg, R) for R in ops]
        _, class_map = class_aggregation(pg, raw_labels)
        class_sizes = Counter(class_map.get(l, l) for l in raw_labels)

    mults = {}
    for ir, row in irreps.items():
        s = 0.0 + 0.0j
        for cl, chi_ir in row.items():
            s += class_sizes.get(cl, 0) * np.conj(chi_ir) * chars.get(cl, 0.0)
        m = s / h
        mults[ir] = int(round(float(np.real(m))))
    return mults

def pretty_gamma(pg: str, mults: Dict[str, int]) -> str:
    order_map = {
        "C2v": ["A1","A2","B1","B2"],
        "C3v": ["A1","A2","E"],
        "C4v": ["A1","A2","B1","B2","E"],
        "C6v": ["A1","A2","B1","B2","E1","E2"],
        "C1": ["A"], "Ci": ["Ag","Au"], "Cs": ["A'","A''"], "Ch": ["A'","A''"],
        "C2": ["A","B"], "C3": ["A","E"], "C4": ["A","B","E"], "C6": ["A","B","E1","E2"],
        "C2h":["Ag","Bg","Au","Bu"],
        "C3h":["A'","A''","E'","E''"],
        "C4h":[f"Γ{a}g" for a in (0,1,2,3)] + [f"Γ{a}u" for a in (0,1,2,3)],
        "D2":["A","B1","B2","B3"],
        "D2h":["Ag","B1g","B2g","B3g","Au","B1u","B2u","B3u"],
        "D3":["A1","A2","E"], "D3d":["A1g","A2g","Eg","A1u","A2u","Eu"], "D3h":["A1'","A2'","E'","A1''","A2''","E''"],
        "D4":["A1","A2","B1","B2","E"], "D2d":["A1","A2","B1","B2","E"], "D4h":["A1g","A2g","B1g","B2g","Eg","A1u","A2u","B1u","B2u","Eu"],
        "D6":["A1","A2","B1","B2","E1","E2"], "D6h":["A1g","A2g","B1g","B2g","E1g","E2g","A1u","A2u","B1u","B2u","E1u","E2u"],
        "T":["A","E","T"], "Td":["A1","A2","E","T1","T2"], "Th":["Ag","Au","Eg","Eu","Tg","Tu"],
        "O":["A1","A2","E","T1","T2"], "Oh":["A1g","A2g","Eg","T1g","T2g","A1u","A2u","Eu","T1u","T2u"],
        "I":["A","T1","T2","G","H"],
        "Ih":["Ag","T1g","T2g","Gg","Hg","Au","T1u","T2u","Gu","Hu"],
    }
    seq = order_map.get(pg, list(CHAR_TABLES[pg]["irreps"].keys()))
    parts = []
    for ir in seq:
        m = mults.get(ir, 0)
        if m > 0:
            parts.append(f"{m}{ir}" if m > 1 else ir)
    return " + ".join(parts) if parts else "0"

# ===================== Abstract-mode helpers =====================
def general_position_chi_3N(classes: List[str], n_atoms: int) -> Dict[str, float]:
    """
    General-position assumption: no atom is fixed by any non-identity operation.
    Then χ_3N(E) = 3N, others 0.
    """
    d = {c: 0.0 for c in classes}
    d["E"] = 3.0 * n_atoms
    return d

def build_raw_labels_from_classes(class_sizes: Dict[str,int]) -> List[str]:
    labels = []
    for cl, sz in class_sizes.items():
        labels.extend([cl]*sz)
    return labels

# ===================== Main =====================
def main():
    ap = argparse.ArgumentParser(
        description="Vibrational irreps from XYZ and a point-group (Schoenflies-like)."
    )
    ap.add_argument("--pg", required=True, help=f"Point group (Schoenflies-like). Supported: {', '.join(sorted(SUPPORTED))}")
    ap.add_argument("--xyz", required=True, help="XYZ file path")
    ap.add_argument("--tol", type=float, default=1e-5, help="Tolerance (Å) for fixing an atom by an operation")
    args = ap.parse_args()

    pg = args.pg
    if pg not in SUPPORTED:
        print(f"Error: point group '{pg}' not supported yet. Supported: {sorted(SUPPORTED)}")
        input("\nPress ENTER to terminate>>\n")
        sys.exit(1)

    coords = load_coords_from_xyz(args.xyz)
    if coords.size == 0:
        print("Error: no coordinates parsed from XYZ.")
        input("\nPress ENTER to terminate>>\n")
        sys.exit(1)

    abstract_mode = pg in ABSTRACT_CLASS_SIZES

    if not abstract_mode:
        # ----- normal mode via pymatgen ops -----
        ops = pg_ops(pg)
        raw_labels = [classify_op_for_table(pg, R) for R in ops]
        classes, class_map = class_aggregation(pg, raw_labels)

        chi_3N_lab = gamma_3N_characters(coords, ops, raw_labels, tol=args.tol)
        chi_T_lab  = gamma_trans_characters(raw_labels, class_map)
        chi_R_lab  = gamma_rot_characters(raw_labels, class_map)

        chi_3N = reduce_to_classes(chi_3N_lab, classes, class_map)
        chi_T  = {c: chi_T_lab.get(c, 0.0) for c in classes}
        chi_R  = {c: chi_R_lab.get(c, 0.0) for c in classes}
        chi_v  = {c: chi_3N[c] - chi_T[c] - chi_R[c] for c in classes}

        mults = decompose(chi_v, pg)

    else:
        # ----- abstract mode (C3h, I, Ih) -----
        classes = CHAR_TABLES[pg]["classes"]
        class_sizes = ABSTRACT_CLASS_SIZES[pg]
        raw_labels = build_raw_labels_from_classes(class_sizes)
        # Γ_T, Γ_R from labels
        # Here class_map is identity (labels already canonical)
        chi_T_lab  = gamma_trans_characters(raw_labels, {c:c for c in classes})
        chi_R_lab  = gamma_rot_characters(raw_labels, {c:c for c in classes})
        chi_T  = {c: chi_T_lab.get(c, 0.0) for c in classes}
        chi_R  = {c: chi_R_lab.get(c, 0.0) for c in classes}
        # Γ_3N: general position assumption
        chi_3N = general_position_chi_3N(classes, coords.shape[0])
        chi_v  = {c: chi_3N[c] - chi_T[c] - chi_R[c] for c in classes}

        mults = decompose(chi_v, pg, class_sizes_override=class_sizes,
                          order_override=sum(class_sizes.values()))

    def row(d): return "  ".join(f"{cl:>7s}:{d[cl]:>6.1f}" for cl in classes)

    print(f"Point group: {pg}  (HM = {to_hm(pg)})  mode={'abstract' if abstract_mode else 'normal'}")
    print("Classes   :", "  ".join(f"{cl:>7s}" for cl in classes))
    print("Γ_3N chars:", row(chi_3N))
    print("Γ_T  chars:", row(chi_T))
    print("Γ_R  chars:", row(chi_R))
    print("Γ_v  chars:", row(chi_v))
    print("\nVibrational irreps:", pretty_gamma(pg, mults))

if __name__ == "__main__":
    main()
    input("\nPress ENTER to terminate>>\n")