#!/usr/bin/env python3
"""Build a serial-crystallography "Table 1" statistics summary from one or
more scale_stream.py partialator run directories (the folders containing
check_hkl.log, cc.log, ccstar.log, Rsplit.log, wilson.log, f2mtz.log and
their .dat files).

Only fills in the rows obtainable before structure refinement (Wavelength
through CC*). Everything from "Reflections used in refinement" onwards
(R-work/R-free, atom counts, Ramachandran, etc.) comes from refinement
itself and isn't available here.

Usage:
  python3 table1.py data/dtpaa_20fs_high --stream_file data/dtpaa_20fs_high_76-77_0.74-scaled.stream
  python3 table1.py data/hewl_nBW data/hewl_LBW --stream_file data/nBW.stream data/LBW.stream -o table1.csv
"""

import argparse
import csv
import os
import re

HEADER = "Reflections measured after indexing"

ROW_ORDER = [
    "Number of crystals",
    "Wavelength",
    "Resolution range",
    "Space group",
    "Unit cell",
    "Total reflections",
    "Unique reflections",
    "Multiplicity",
    "Completeness (%)",
    "Mean I/sigma(I)",
    "Wilson B-factor",
    "R-split",
    "CC1/2",
    "CC*",
]


def _read(path):
    with open(path) as f:
        return f.read()


def _search(text, pattern, cast=float):
    match = re.search(pattern, text)
    return cast(match.group(1)) if match else None


def _last_shell_row(dat_path):
    with open(dat_path) as f:
        lines = [line for line in f if line.strip()]
    return lines[-1].split()


def resolution_range(cc_log_text):
    return (
        _search(cc_log_text, r"Accepted resolution range:.*\(([\d.]+) to [\d.]+ Angstroms\)"),
        _search(cc_log_text, r"Accepted resolution range:.*\([\d.]+ to ([\d.]+) Angstroms\)"),
    )


def last_shell_resolution(mult_dat_path):
    fields = _last_shell_row(mult_dat_path)
    min_1nm, max_1nm = float(fields[-2]), float(fields[-1])
    return 10.0 / min_1nm, 10.0 / max_1nm


def spacegroup_and_cell(f2mtz_log_text):
    cell_match = re.search(r"Data line--- CELL\s+([\d. ]+)", f2mtz_log_text)
    cell = [float(v) for v in cell_match.group(1).split()] if cell_match else None
    sg_match = re.search(r"Space Group\s*=\s*\d+\s*'([^']+)'", f2mtz_log_text)
    spacegroup = sg_match.group(1) if sg_match else None
    return spacegroup, cell


def wavelength_from_stream(stream_path):
    with open(stream_path) as f:
        for line in f:
            match = re.match(r"photon_energy\s*=\s*([\d.]+)\s*eV", line.strip())
            if match:
                energy_ev = float(match.group(1))
                return 12398.4193 / energy_ev
    return None


def count_crystals(stream_path):
    with open(stream_path) as f:
        return sum(1 for line in f if line.strip() == HEADER)


def gather_stats(proc_dir, stream_file=None):
    check_hkl_log = _read(os.path.join(proc_dir, "check_hkl.log"))
    cc_log = _read(os.path.join(proc_dir, "cc.log"))
    ccstar_log = _read(os.path.join(proc_dir, "ccstar.log"))
    rsplit_log = _read(os.path.join(proc_dir, "Rsplit.log"))
    wilson_log = _read(os.path.join(proc_dir, "wilson.log"))
    f2mtz_log = _read(os.path.join(proc_dir, "f2mtz.log"))

    mult_dat = os.path.join(proc_dir, "mult.dat")
    cc_dat = os.path.join(proc_dir, "cc.dat")
    ccstar_dat = os.path.join(proc_dir, "ccstar.dat")
    rsplit_dat = os.path.join(proc_dir, "Rsplit.dat")

    low_res, high_res = resolution_range(cc_log)
    shell_low, shell_high = last_shell_resolution(mult_dat)
    spacegroup, cell = spacegroup_and_cell(f2mtz_log)

    mult_row = _last_shell_row(mult_dat)
    shell_unique = int(mult_row[1])
    shell_comp = float(mult_row[3])
    shell_mult = float(mult_row[5])
    shell_snr = float(mult_row[6])

    shell_cc = float(_last_shell_row(cc_dat)[1])
    shell_ccstar = float(_last_shell_row(ccstar_dat)[1])
    shell_rsplit = float(_last_shell_row(rsplit_dat)[1])

    total_measurements = _search(check_hkl_log, r"(\d+) measurements in total", int)
    total_unique = _search(check_hkl_log, r"(\d+) reflections in total\.", int)
    overall_mult = _search(check_hkl_log, r"Overall redundancy = ([\d.]+)")
    overall_comp = _search(check_hkl_log, r"Overall completeness = ([\d.]+)")
    overall_snr = _search(check_hkl_log, r"Overall\s<snr>\s=\s([\d.]+)")

    wilson_b = _search(wilson_log, r"B = ([\d.]+) A")
    overall_rsplit = _search(rsplit_log, r"Overall Rsplit = ([\d.]+)")
    overall_cc = _search(cc_log, r"Overall CC = ([\d.]+)")
    overall_ccstar = _search(ccstar_log, r"Overall CC\* = ([\d.]+)")

    wavelength = wavelength_from_stream(stream_file) if stream_file else None
    n_crystals = count_crystals(stream_file) if stream_file else None

    return {
        "Wavelength": f"{wavelength:.2f}" if wavelength else "N/A",
        "Resolution range": f"{low_res:.2f} - {high_res:.2f} ({shell_low:.2f} - {shell_high:.2f})",
        "Space group": spacegroup or "N/A",
        "Unit cell": " ".join(f"{v:.2f}" for v in cell) if cell else "N/A",
        "Number of crystals": n_crystals if n_crystals is not None else "N/A",
        "Total reflections": total_measurements if total_measurements is not None else "N/A",
        "Unique reflections": f"{total_unique} ({shell_unique})" if total_unique is not None else "N/A",
        "Multiplicity": f"{overall_mult:.2f} ({shell_mult:.2f})" if overall_mult is not None else "N/A",
        "Completeness (%)": f"{overall_comp:.2f} ({shell_comp:.2f})" if overall_comp is not None else "N/A",
        "Mean I/sigma(I)": f"{overall_snr:.2f} ({shell_snr:.2f})" if overall_snr is not None else "N/A",
        "Wilson B-factor": f"{wilson_b:.2f}" if wilson_b is not None else "N/A",
        "R-split": f"{overall_rsplit:.2f} ({shell_rsplit:.2f})" if overall_rsplit is not None else "N/A",
        "CC1/2": f"{overall_cc:.2f} ({shell_cc:.2f})" if overall_cc is not None else "N/A",
        "CC*": f"{overall_ccstar:.2f} ({shell_ccstar:.2f})" if overall_ccstar is not None else "N/A",
    }


def main():
    parser = argparse.ArgumentParser(
        description="Build a Table 1 statistics summary from scale_stream.py "
        "partialator run directories."
    )
    parser.add_argument(
        "proc_dirs", nargs="+",
        help="One or more partialator run directories, e.g. data/dtpaa_20fs_high",
    )
    parser.add_argument(
        "--stream_file", nargs="+",
        help="The .stream file for each proc_dir, in the same order, used to read the "
        "wavelength (from photon_energy) and count the number of crystals. Omit a "
        "directory's by passing '' in its place. If not given at all, both are "
        "left as N/A.",
    )
    parser.add_argument(
        "-o", "--output", help="Write the table as a csv file to this path."
    )
    args = parser.parse_args()

    if args.stream_file and len(args.stream_file) != len(args.proc_dirs):
        parser.error(
            f"--stream_file needs one entry per proc_dir "
            f"({len(args.proc_dirs)} given), got {len(args.stream_file)}"
        )
    stream_files = args.stream_file or [None] * len(args.proc_dirs)

    columns = {}
    for proc_dir, stream_file in zip(args.proc_dirs, stream_files):
        label = os.path.basename(os.path.normpath(proc_dir))
        columns[label] = gather_stats(proc_dir, stream_file or None)

    rows = [[""] + list(columns.keys())]
    for row in ROW_ORDER:
        rows.append([row] + [columns[label][row] for label in columns])

    print("\n".join(", ".join(str(v) for v in row) for row in rows))

    if args.output:
        with open(args.output, "w", newline="") as f:
            csv.writer(f).writerows(rows)
        print(f"\nWrote table to {args.output}")


if __name__ == "__main__":
    main()