#!/usr/bin/env python3
"""Decompose fast-mode-versus-standard HDR error by scalar RGB sample value.

The unweighted whole-image relative L2 is dominated by a small number of very
bright PSF cores, so it is not a good proxy for how the dense faint-star
texture is reproduced.  FITS RGB channels are flattened into scalar samples.
This tool reports, for each fast-mode FITS image:

  * the cumulative fraction of ||fast-base||^2 contributed by the largest |d|
    RGB channel samples, alongside their share of the baseline linear-RGB sum;
  * local relative RMS, signed bias, and share of base energy and linear-RGB
    sum for a disjoint partition of base-value ranges.

Usage:
    python3 scripts/fast_mode_error_decomposition.py BASE.fits FAST.fits [...]
"""
import sys

import numpy as np


def read_primary_fits(path):
    with open(path, 'rb') as stream:
        header = b''
        while True:
            block = stream.read(2880)
            if not block:
                raise ValueError(f'{path}: truncated FITS header')
            header += block
            if any(block[i:i + 8] == b'END     ' for i in range(0, 2880, 80)):
                break
        cards = {}
        for i in range(0, len(header), 80):
            card = header[i:i + 80].decode('ascii', 'replace')
            key = card[:8].strip()
            if key in ('NAXIS1', 'NAXIS2', 'NAXIS3', 'BITPIX', 'NAXIS'):
                cards[key] = int(card[10:30])
            if key == 'END':
                break
        stream.seek(len(header))
        count = cards['NAXIS1'] * cards.get('NAXIS2', 1) * cards.get('NAXIS3', 1)
        dtype = '>f4' if cards['BITPIX'] == -32 else '>f8'
        return np.frombuffer(stream.read(count * np.dtype(dtype).itemsize),
                             dtype=dtype).astype(np.float64)


def decompose(name, base, fast, ranks=(1e-5, 1e-4, 1e-3, 1e-2, 1e-1)):
    diff = fast - base
    abs_diff = np.abs(diff)
    base_l2_sq = float(np.dot(base, base))
    diff_l2_sq = float(np.dot(abs_diff, abs_diff))
    base_rgb_sum = float(base.sum())
    print(f'\n===== {name} =====  relative_L2={np.sqrt(diff_l2_sq / base_l2_sq) * 100:.4f}%')
    order = np.argsort(abs_diff)[::-1]
    print('top |d| RGB samples:  sq-error energy   base energy   linear-RGB sum   sample share')
    for frac in ranks:
        k = max(1, int(frac * base.size))
        idx = order[:k]
        print(f'  top {frac * 100:8.4f}%: {np.dot(abs_diff[idx], abs_diff[idx]) / diff_l2_sq * 100:7.3f}%'
              f'  {np.dot(base[idx], base[idx]) / base_l2_sq * 100:9.3f}%'
              f'  {base[idx].sum() / base_rgb_sum * 100:8.3f}%'
              f'  {k / base.size * 100:9.4f}%')
    nonzero = base[base > 0]
    quantiles = np.quantile(nonzero, [0.5, 0.9, 0.99, 0.999])
    edges = np.unique(np.concatenate(
        ([-np.inf, 1e-6, 1e-3, 1e-1], quantiles, [np.inf])))
    print('base-value range (disjoint):'
          '  rel_RMS   signed bias   base energy %   linear-RGB sum %   samples')
    for low, high in zip(edges[:-1], edges[1:]):
        mask = (base > low) & (base <= high)
        if not mask.any():
            continue
        bin_diff = diff[mask]
        bin_base = base[mask]
        rel = np.sqrt(np.dot(bin_diff, bin_diff) / np.dot(bin_base, bin_base))
        bias = bin_diff.sum() / bin_base.sum()
        lo = '-inf' if low == -np.inf else f'{low:.3g}'
        hi = 'inf' if high == np.inf else f'{high:.3g}'
        print(f'  [{lo},{hi}): {rel * 100:9.4f}%  {bias * 100:+10.4f}%'
              f'  {np.dot(bin_base, bin_base) / base_l2_sq * 100:8.3f}%'
              f'  {bin_base.sum() / base_rgb_sum * 100:7.3f}%  {int(mask.sum())}')


def main(argv):
    if len(argv) < 3:
        print(__doc__)
        return 2
    base = read_primary_fits(argv[1])
    for path in argv[2:]:
        fast = read_primary_fits(path)
        if fast.shape != base.shape:
            raise ValueError(f'{path}: shape {fast.shape} != base {base.shape}')
        decompose(path.rsplit('/', 1)[-1], base, fast)
    return 0


if __name__ == '__main__':
    sys.exit(main(sys.argv))