#!/usr/bin/env python3

import argparse
import sys
import numpy as np
import gemmi


def read_mtz_amplitudes(mtz_path, f_label):
    """
    Reads an MTZ file and returns a dictionary:
        {(h,k,l): amplitude}
    """
    mtz = gemmi.read_mtz_file(mtz_path)
    mtz.setup_batched()

    try:
        col = mtz.column_with_label(f_label)
    except RuntimeError:
        sys.exit(f"ERROR: Column '{f_label}' not found in {mtz_path}")

    hkl_dict = {}

    for row in mtz:
        h, k, l = int(row[0]), int(row[1]), int(row[2])
        fval = row[col.idx]

        if not np.isnan(fval):
            hkl_dict[(h, k, l)] = float(fval)

    return hkl_dict


def match_reflections(dict1, dict2):
    """
    Returns matched amplitude arrays for common HKLs.
    """
    common_hkls = set(dict1.keys()) & set(dict2.keys())

    if len(common_hkls) == 0:
        sys.exit("ERROR: No matching HKLs between files.")

    f1 = np.array([dict1[hkl] for hkl in common_hkls])
    f2 = np.array([dict2[hkl] for hkl in common_hkls])

    return f1, f2, len(common_hkls)


def compute_riso(f1, f2):
    """
    Computes Riso:
        Riso = sum | |F1| - |F2| | / sum |F1|
    """
    numerator = np.sum(np.abs(np.abs(f1) - np.abs(f2)))
    denominator = np.sum(np.abs(f1))

    if denominator == 0:
        sys.exit("ERROR: Denominator is zero.")

    return numerator / denominator


def main():
    parser = argparse.ArgumentParser(
        description="Compute Riso between two MTZ files."
    )

    parser.add_argument("mtz1", help="Reference MTZ file")
    parser.add_argument("mtz2", help="Second MTZ file")
    parser.add_argument("--f1", required=True,
                        help="Amplitude column label in first MTZ")
    parser.add_argument("--f2", required=True,
                        help="Amplitude column label in second MTZ")

    args = parser.parse_args()

    print("Reading MTZ files...")
    dict1 = read_mtz_amplitudes(args.mtz1, args.f1)
    dict2 = read_mtz_amplitudes(args.mtz2, args.f2)

    print("Matching reflections...")
    f1, f2, nmatch = match_reflections(dict1, dict2)

    print(f"Matched reflections: {nmatch}")

    riso_value = compute_riso(f1, f2)

    print(f"\nRiso = {riso_value:.6f}")


if __name__ == "__main__":
    main()