#!/usr/bin/env python3
"""lhe_reader.py — a dependency-free, streaming reader for Les Houches Event files.

Handles LHEF 1.0 (hep-ph/0609017) and the LHEF 3.0 additions (arXiv:1405.1067):
<init>, <event> blocks, per-event <rwgt> weights, <initrwgt> weight definitions,
event attributes, and gzip-compressed files. It deliberately does NOT use an XML
parser: real LHE files carry unescaped '<', '&' and raw generator banners in the
<header>, so a strict parser will choke. A line-oriented state machine is what
every production reader (Pythia's LHAup, Rivet, pylhe) actually does.

Usage:
    from lhe_reader import read_lhe, check_event
    with open("events.lhe") as f:          # or gzip.open(..., "rt")
        init, events = read_lhe(f)
        for ev in events:
            print(ev.weight, [p.pid for p in ev.particles if p.status == 1])
"""
from __future__ import annotations
import gzip, math, re, sys
from dataclasses import dataclass, field
from typing import Dict, Iterator, List, Optional, TextIO, Tuple


# ------------------------------------------------------------------ records
@dataclass
class Particle:
    pid: int            # IDUP     PDG Monte Carlo numbering scheme
    status: int         # ISTUP    -1 in, +1 out, 2 resonance, -2 spacelike, 3 doc, -9 beam
    mother1: int        # MOTHUP(1)  1-based index into the event's particle list (0 = none)
    mother2: int        # MOTHUP(2)
    color: int          # ICOLUP(1)  colour line label (0 = none)
    anticolor: int      # ICOLUP(2)  anticolour line label
    px: float; py: float; pz: float; e: float; m: float   # PUP(1..5), GeV
    lifetime: float     # VTIMUP   c*tau, mm
    spin: float         # SPINUP   helicity / cos(angle to mother); 9 = unknown

    @property
    def pt(self) -> float: return math.hypot(self.px, self.py)
    @property
    def p(self) -> float: return math.sqrt(self.px**2 + self.py**2 + self.pz**2)
    @property
    def eta(self) -> float:
        p = self.p
        return 0.5 * math.log((p + self.pz) / (p - self.pz)) if p > abs(self.pz) else math.copysign(float("inf"), self.pz)
    @property
    def phi(self) -> float: return math.atan2(self.py, self.px)
    @property
    def mass_from_p4(self) -> float:
        m2 = self.e**2 - self.px**2 - self.py**2 - self.pz**2
        return math.copysign(math.sqrt(abs(m2)), m2)


@dataclass
class Event:
    nup: int            # NUP      number of particle entries
    process: int        # IDPRUP   which LPRUP process this event belongs to
    weight: float       # XWGTUP   event weight (meaning fixed by IDWTUP)
    scale: float        # SCALUP   factorisation/shower starting scale, GeV
    alpha_qed: float    # AQEDUP
    alpha_qcd: float    # AQCDUP
    particles: List[Particle] = field(default_factory=list)
    weights: Dict[str, float] = field(default_factory=dict)   # LHEF3 <rwgt><wgt id=...>
    attributes: Dict[str, str] = field(default_factory=dict)  # attributes on the <event> tag
    extra: List[str] = field(default_factory=list)            # any other lines inside <event>

    def final_state(self) -> List[Particle]:
        return [p for p in self.particles if p.status == 1]

    def incoming(self) -> List[Particle]:
        return [p for p in self.particles if p.status == -1]

    def p4_sum(self, parts) -> Tuple[float, float, float, float]:
        return (sum(p.px for p in parts), sum(p.py for p in parts),
                sum(p.pz for p in parts), sum(p.e for p in parts))


@dataclass
class ProcessInfo:
    xsec: float         # XSECUP  pb
    xsec_err: float     # XERRUP  pb
    xmax: float         # XMAXUP  maximum weight
    lprup: int          # LPRUP   process label used by IDPRUP


@dataclass
class Init:
    beam_pids: Tuple[int, int]        # IDBMUP
    beam_energies: Tuple[float, float] # EBMUP, GeV
    pdf_groups: Tuple[int, int]        # PDFGUP
    pdf_sets: Tuple[int, int]          # PDFSUP  (LHAPDF ids)
    weight_strategy: int               # IDWTUP  ±1..±4 (see docs)
    processes: List[ProcessInfo]       # NPRUP entries
    weight_names: Dict[str, str] = field(default_factory=dict)   # LHEF3 <initrwgt> id -> description
    extra: List[str] = field(default_factory=list)
    header: str = ""                   # raw text of the <header> block


# ------------------------------------------------------------------ parsing
_TAG_ATTR = re.compile(r'(\w+)\s*=\s*"([^"]*)"')


def _strip_comment(line: str) -> str:
    """Everything after '#' is a comment in the numeric blocks."""
    i = line.find("#")
    return line if i < 0 else line[:i]


def _parse_init(lines: List[str], weight_names: Dict[str, str], header: str) -> Init:
    numeric = [l for l in (_strip_comment(x).strip() for x in lines) if l]
    f = numeric[0].split()
    if len(f) < 10:
        raise ValueError(f"malformed <init> first line: {numeric[0]!r}")
    nprup = int(f[9])
    procs = []
    for l in numeric[1:1 + nprup]:
        x = l.split()
        procs.append(ProcessInfo(float(x[0]), float(x[1]), float(x[2]), int(x[3])))
    return Init(beam_pids=(int(f[0]), int(f[1])),
                beam_energies=(float(f[2]), float(f[3])),
                pdf_groups=(int(f[4]), int(f[5])),
                pdf_sets=(int(f[6]), int(f[7])),
                weight_strategy=int(f[8]),
                processes=procs,
                weight_names=weight_names,
                extra=numeric[1 + nprup:],
                header=header)


def _parse_event(lines: List[str], attrs: Dict[str, str]) -> Event:
    """lines = everything between <event ...> and </event>."""
    # 1. peel off LHEF3 <rwgt> ... </rwgt> and remember any other tagged blocks verbatim
    weights: Dict[str, float] = {}
    body: List[str] = []
    extra: List[str] = []
    in_rwgt = False
    for raw in lines:
        s = raw.strip()
        if not s:
            continue
        if s.startswith("<rwgt"):
            in_rwgt = True; continue
        if s.startswith("</rwgt"):
            in_rwgt = False; continue
        if in_rwgt:
            m = re.match(r"<wgt\s+id\s*=\s*['\"]([^'\"]+)['\"]\s*>\s*([-+0-9.eEdD]+)\s*</wgt>", s)
            if m:
                weights[m.group(1)] = float(m.group(2).replace("D", "E").replace("d", "e"))
            continue
        if s.startswith("<"):            # <mgrwt>, <scales>, <clustering>, ... keep, don't parse
            extra.append(s); continue
        body.append(_strip_comment(s))
    body = [b for b in body if b.strip()]
    if not body:
        raise ValueError("empty <event> block")
    h = body[0].split()
    nup = int(h[0])
    ev = Event(nup=nup, process=int(h[1]), weight=float(h[2]), scale=float(h[3]),
               alpha_qed=float(h[4]), alpha_qcd=float(h[5]), weights=weights,
               attributes=attrs, extra=extra)
    if len(body) < 1 + nup:
        raise ValueError(f"event declares NUP={nup} but has {len(body)-1} particle lines")
    for l in body[1:1 + nup]:
        x = l.replace("D", "E").replace("d", "e").split()   # Fortran 1.0D+00 style
        ev.particles.append(Particle(
            pid=int(x[0]), status=int(x[1]), mother1=int(x[2]), mother2=int(x[3]),
            color=int(x[4]), anticolor=int(x[5]),
            px=float(x[6]), py=float(x[7]), pz=float(x[8]), e=float(x[9]), m=float(x[10]),
            lifetime=float(x[11]), spin=float(x[12])))
    return ev


def _parse_initrwgt(lines: List[str]) -> Dict[str, str]:
    names: Dict[str, str] = {}
    group = ""
    for s in (l.strip() for l in lines):
        if s.startswith("<weightgroup"):
            m = re.search(r'name\s*=\s*["\']([^"\']*)["\']', s)
            group = m.group(1) if m else ""
        m = re.match(r"<weight\s+id\s*=\s*['\"]([^'\"]+)['\"][^>]*>(.*?)</weight>", s)
        if m:
            names[m.group(1)] = (group + " : " if group else "") + m.group(2).strip()
    return names


def read_lhe(f: TextIO) -> Tuple[Init, Iterator[Event]]:
    """Return (init, iterator over events). The iterator streams: a 20 GB file is fine."""
    header_lines: List[str] = []
    init_lines: List[str] = []
    initrwgt_lines: List[str] = []
    state = "top"
    line_no = 0
    # ---- phase 1: read up to the end of <init> (header + init are always first)
    for raw in f:
        line_no += 1
        s = raw.strip()
        if state == "top":
            if s.startswith("<header"):
                state = "header"
            elif s.startswith("<init"):
                state = "init"
        elif state == "header":
            if s.startswith("</header"):
                state = "top"
            else:
                header_lines.append(raw)
                if s.startswith("<initrwgt"):
                    state = "initrwgt"
        elif state == "initrwgt":
            header_lines.append(raw)
            if s.startswith("</initrwgt"):
                state = "header"
            else:
                initrwgt_lines.append(raw)
        elif state == "init":
            if s.startswith("</init"):
                break
            init_lines.append(raw)
    else:
        raise ValueError("no <init> block found — is this an LHE file?")
    init = _parse_init(init_lines, _parse_initrwgt(initrwgt_lines), "".join(header_lines))

    # ---- phase 2: a generator that yields events one at a time
    def events() -> Iterator[Event]:
        buf: List[str] = []
        attrs: Dict[str, str] = {}
        inside = False
        for raw in f:
            s = raw.strip()
            if not inside:
                if s.startswith("<event"):
                    inside = True
                    attrs = dict(_TAG_ATTR.findall(s))
                    buf = []
                # anything else at top level (blank lines, </LesHouchesEvents>) is ignored
            else:
                if s.startswith("</event"):
                    inside = False
                    yield _parse_event(buf, attrs)
                else:
                    buf.append(raw)
    return init, events()


def open_lhe(path: str) -> TextIO:
    """Transparently open plain or gzipped files in text mode."""
    return gzip.open(path, "rt") if path.endswith(".gz") else open(path, "r")


# ------------------------------------------------------------------ checks
def check_event(ev: Event, tol: float = 1e-3) -> List[str]:
    """Sanity checks that catch most broken files. Returns a list of complaints (empty = fine)."""
    out: List[str] = []
    n = len(ev.particles)
    if n != ev.nup:
        out.append(f"NUP={ev.nup} but {n} particles read")
    for i, p in enumerate(ev.particles, 1):
        if p.status not in (-9, -2, -1, 1, 2, 3):
            out.append(f"particle {i}: unknown status {p.status}")
        for mo in (p.mother1, p.mother2):
            if mo < 0 or mo > n:
                out.append(f"particle {i}: mother index {mo} out of range 1..{n}")
            elif mo >= i and mo != 0:
                out.append(f"particle {i}: mother {mo} does not precede it")
        # on-shell consistency: the stored mass should match E^2 - p^2
        m_calc = p.mass_from_p4
        if abs(m_calc - p.m) > tol * max(1.0, p.e):
            out.append(f"particle {i} (pid {p.pid}): m={p.m:.6g} but sqrt(E^2-p^2)={m_calc:.6g}")
    # four-momentum conservation: incoming == outgoing (status 1 only; resonances are intermediate)
    pin, pout = ev.p4_sum(ev.incoming()), ev.p4_sum(ev.final_state())
    for k, (a, b) in zip("px py pz E".split(), zip(pin, pout)):
        if abs(a - b) > tol * max(1.0, abs(a), abs(b)):
            out.append(f"{k} not conserved: in {a:.6g}, out {b:.6g}")
    # colour flow: every colour label must appear once as colour and once as anticolour
    # among the 'external' legs (incoming legs count with colour/anticolour swapped)
    seen: Dict[int, int] = {}
    for p in ev.particles:
        if p.status not in (-1, 1):
            continue
        c, a = (p.anticolor, p.color) if p.status == -1 else (p.color, p.anticolor)
        if c: seen[c] = seen.get(c, 0) + 1
        if a: seen[a] = seen.get(a, 0) - 1
    for lbl, bal in seen.items():
        if bal != 0:
            out.append(f"colour line {lbl} unbalanced ({bal:+d})")
    return out


def cross_section(init: Init, weights: List[float]) -> Tuple[float, float]:
    """Total cross section (pb) implied by a sample, and its MC statistical error.

    IDWTUP = ±3 : unweighted, |w| = 1 -> sigma is XSECUP from <init>; only the sign matters.
    IDWTUP = ±4 : weighted, sum(w)/N = sigma (pb) -> sigma = mean weight.
    IDWTUP = ±1,±2 : weights are only meaningful relative to XMAXUP; use XSECUP.
    """
    n = len(weights)
    strat = abs(init.weight_strategy)
    if strat == 4:
        mean = sum(weights) / n
        var = sum((w - mean) ** 2 for w in weights) / (n - 1) if n > 1 else 0.0
        return mean, math.sqrt(var / n)
    sigma = sum(p.xsec for p in init.processes)
    if strat == 3:
        # sign-only weights: sigma * <sign> with the same fractional error
        mean_sign = sum(math.copysign(1.0, w) for w in weights) / n
        return sigma * mean_sign, sigma / math.sqrt(n)
    return sigma, math.sqrt(sum(p.xsec_err ** 2 for p in init.processes))


# ------------------------------------------------------------------ CLI
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__); sys.exit(0)
    with open_lhe(sys.argv[1]) as fh:
        init, evs = read_lhe(fh)
        print(f"beams {init.beam_pids} at {init.beam_energies} GeV, PDF sets {init.pdf_sets}, IDWTUP={init.weight_strategy}")
        for p in init.processes:
            print(f"  process {p.lprup}: sigma = {p.xsec:.6g} +- {p.xsec_err:.2g} pb")
        if init.weight_names:
            print(f"  {len(init.weight_names)} named weights, e.g. {next(iter(init.weight_names.items()))}")
        n = neg = bad = 0
        ws: List[float] = []
        for ev in evs:
            n += 1; ws.append(ev.weight); neg += ev.weight < 0
            problems = check_event(ev)
            if problems and bad < 5:
                print(f"event {n}: " + "; ".join(problems))
            bad += bool(problems)
        sig, err = cross_section(init, ws)
        print(f"{n} events, {neg} with negative weight ({100*neg/max(n,1):.1f}%), {bad} failed checks")
        print(f"sigma from weights = {sig:.6g} +- {err:.2g} pb")