#!python
"""Print summary statistics for a FastSpecFit VAC.

Examples:

vac-statistics --specprod=iron --vacdir=$DESI_ROOT/public/dr1/vac/dr1/fastspecfit/iron/v3.0


"""
import os
import argparse
from glob import glob
import numpy as np
import fitsio

from fastspecfit.logger import getFastspecLogger
log = getFastspecLogger()

_SURVEY_ORDER = {'cmx': 0, 'sv1': 1, 'sv2': 2, 'sv3': 3, 'special': 4, 'main': 5}
_PROGRAM_ORDER = {'backup': 0, 'bright': 1, 'dark': 2, 'other': 3}

_PACKAGES = ['python', 'numpy', 'scipy', 'astropy', 'yaml', 'matplotlib',
             'fitsio', 'mpi4py', 'healpy', 'desiutil', 'desispec',
             'desitarget', 'speclite', 'fastspecfit']

_ENV_VARS = ['DESI_ROOT', 'DUST_DIR', 'FPHOTO_DIR', 'FTEMPLATES_DIR',
             'FTEMPLATES_FILE', 'FPHOTO_FILE', 'EMLINES_FILE',
             'CONSTRAINTS_FILE']


def _parse_survey_program(filepath, specprod, script):
    base = os.path.basename(filepath).replace('.fits', '')
    rest = base[len(f'{script}-{specprod}'):].lstrip('-')
    parts = rest.split('-') if rest else []
    survey = parts[0] if parts else ''
    program = parts[1] if len(parts) > 1 else ''
    return survey, program


def _catalog_sort_key(filepath, specprod, script):
    survey, program = _parse_survey_program(filepath, specprod, script)
    return (_SURVEY_ORDER.get(survey, 99), _PROGRAM_ORDER.get(program, 99))


def _survey_groups(files, specprod, script):
    """Yield (survey, start, end) index ranges for contiguous survey groups."""
    surveys = [_parse_survey_program(f, specprod, script)[0] for f in files]
    i = 0
    while i < len(surveys):
        survey = surveys[i]
        j = i
        while j < len(surveys) and surveys[j] == survey:
            j += 1
        yield survey, i, j
        i = j


def _format_size(sz):
    if sz < 1024**2:
        return '{:.3g} KB'.format(sz / 1024)
    elif sz < 1024**3:
        return '{:.3g} MB'.format(sz / 1024**2)
    else:
        return '{:.3g} GB'.format(sz / 1024**3)


def gather_catalogs(vacdir, specprod, script):
    catdir = os.path.join(vacdir, 'catalogs')
    files = (
        sorted(glob(os.path.join(catdir, f'{script}-{specprod}-*-*.fits'))) +
        glob(os.path.join(catdir, f'{script}-{specprod}.fits')) +
        glob(os.path.join(catdir, f'{script}-{specprod}-special.fits')) +
        glob(os.path.join(catdir, f'{script}-{specprod}-sv.fits')) +
        glob(os.path.join(catdir, f'{script}-{specprod}-main.fits'))
    )
    files = sorted(files, key=lambda f: _catalog_sort_key(f, specprod, script))
    return np.array(files)


def fileinfo(filelist):

    headers = ['File Name', 'File Size', 'Number of Targets']

    nrows_raw, szs_raw = [], []
    for onefile in filelist:
        szs_raw.append(os.stat(onefile).st_size)
        nrows_raw.append(fitsio.FITS(onefile)[1].get_nrows())
    szs_raw = np.array(szs_raw)
    nrows_raw = np.array(nrows_raw)

    lines = []
    widths = np.array([len(header) for header in headers])
    for onefile, nrow, sz in zip(filelist, nrows_raw, szs_raw):
        basefile = os.path.basename(onefile)
        sz_str = _format_size(sz)
        nrow_str = '{:,d}'.format(nrow)

        oneline = []
        for ii, dd in enumerate([basefile, sz_str, nrow_str]):
            if len(dd) > widths[ii]:
                widths[ii] = len(dd)
            oneline.append(dd)

        lines.append(oneline)

    return widths, lines, headers, nrows_raw, szs_raw


def build_table(widths, lines, headers, insert_after=None):

    if insert_after is None:
        insert_after = {}

    def print_rule(char):
        for ii, width in enumerate(widths):
            end = '\n' if ii == len(widths) - 1 else ' '
            print(char * width, end=end)

    def print_line(line):
        for ii, (width, dd) in enumerate(zip(widths, line)):
            end = '\n' if ii == len(widths) - 1 else ' '
            print(str.ljust(dd, width), end=end)

    print_rule('=')
    print_line(headers)
    print_rule('=')

    for i, line in enumerate(lines):
        print_line(line)
        if i in insert_after:
            print_line(insert_after[i])

    print_rule('=')


def print_file_statistics(files, specprod, script):
    """Print one RST table per survey for file sizes and target counts."""
    for survey, i, j in _survey_groups(files, specprod, script):
        survey_files = files[i:j]
        widths, lines, headers, nrows_raw, szs_raw = fileinfo(survey_files)

        insert_after = {}
        if j - i > 1:
            totals_line = [
                f'Total ({survey})',
                _format_size(int(np.sum(szs_raw))),
                '{:,d}'.format(int(np.sum(nrows_raw))),
            ]
            for ii, val in enumerate(totals_line):
                if len(val) > widths[ii]:
                    widths[ii] = len(val)
            insert_after[j - i - 1] = totals_line

        build_table(widths, lines, headers, insert_after=insert_after)
        print()


def print_qso_statistics(files, specprod, script):
    """Print one RST table per survey for QSO redshift corrections."""
    for survey, i, j in _survey_groups(files, specprod, script):
        survey_files = files[i:j]

        headers = ['Catalog', 'Number of Objects', 'Number with Corrected Redshifts']
        widths = np.array([len(header) for header in headers])
        lines = []
        nrows_raw, nznews_raw = [], []

        for fastfile in survey_files:
            log.debug(f'Reading {fastfile}')
            fast = fitsio.read(fastfile, ext='METADATA', columns=['Z', 'Z_RR'])
            nrow = len(fast)
            nznew = int(np.sum(fast['Z'] != fast['Z_RR']))
            nrows_raw.append(nrow)
            nznews_raw.append(nznew)

            basefile = os.path.basename(fastfile)
            nrow_str = '{:,d}'.format(nrow)
            nznew_str = '{:,d}'.format(nznew)

            for ii, dd in enumerate([basefile, nrow_str, nznew_str]):
                if len(dd) > widths[ii]:
                    widths[ii] = len(dd)
            lines.append([basefile, nrow_str, nznew_str])

        nrows_raw = np.array(nrows_raw)
        nznews_raw = np.array(nznews_raw)

        insert_after = {}
        if j - i > 1:
            totals_line = [
                f'Total ({survey})',
                '{:,d}'.format(int(np.sum(nrows_raw))),
                '{:,d}'.format(int(np.sum(nznews_raw))),
            ]
            for ii, val in enumerate(totals_line):
                if len(val) > widths[ii]:
                    widths[ii] = len(val)
            insert_after[j - i - 1] = totals_line

        build_table(widths, lines, headers, insert_after=insert_after)
        print()


def gather_pixel_files(vacdir, script, sample=True):
    """Return per-pixel catalog files under vacdir/healpix/ and vacdir/spectra/.

    If sample=True (default), descend one directory at a time and return the
    first file found in each (survey, program) combination — typically a few
    dozen files read in under a second.  If sample=False, return every pixel
    file (slow for large specprods).
    """
    ext = '.fits.gz' if script == 'fastspec' else '.fits'
    files = []
    for subdir in ('healpix', 'spectra'):
        pixel_dir = os.path.join(vacdir, subdir)
        if not os.path.isdir(pixel_dir):
            continue
        if not sample:
            pattern = os.path.join(pixel_dir, '*', '*', '*', '*', f'{script}-*{ext}')
            files.extend(glob(pattern))
            continue
        for survey_dir in sorted(glob(os.path.join(pixel_dir, '*'))):
            for program_dir in sorted(glob(os.path.join(survey_dir, '*'))):
                try:
                    pixgroup = sorted(os.listdir(program_dir))[0]
                    pixel = sorted(os.listdir(os.path.join(program_dir, pixgroup)))[0]
                    leaf = os.path.join(program_dir, pixgroup, pixel)
                    match = next((os.path.join(leaf, f) for f in sorted(os.listdir(leaf))
                                  if f.startswith(script) and f.endswith(ext)), None)
                    if match:
                        files.append(match)
                except (OSError, IndexError):
                    continue
    return sorted(files)


def print_code_versions(vacdir, all_files, sample=True):
    """Print software versions table collected from per-pixel catalog headers.

    When sample=True (default), reads one file per survey/program.
    When sample=False, reads every per-pixel file (slow for large specprods).
    """
    from desiutil.depend import Dependencies

    pkg_versions = {pkg: set() for pkg in _PACKAGES}
    nread = 0

    for script, merged_files in all_files.items():
        source_files = gather_pixel_files(vacdir, script, sample=sample)
        if not source_files:
            log.warning(f'No per-pixel files found for {script}; falling back to merged catalogs.')
            source_files = list(merged_files)
        log.info(f'Reading code versions from {len(source_files)} {script} files.')
        for filepath in source_files:
            try:
                hdr = fitsio.read_header(filepath, ext=0)
                for pkg, ver in Dependencies(hdr).items():
                    if pkg in pkg_versions:
                        pkg_versions[pkg].add(str(ver))
                nread += 1
            except Exception as e:
                log.warning(f'Could not read code versions from {filepath}: {e}')

    if nread == 0:
        log.warning('No code versions could be read.')
        return

    headers = ['Software Package', 'Version(s)']
    widths = np.array([len(h) for h in headers])
    lines = []
    for pkg in _PACKAGES:
        versions = pkg_versions[pkg]
        if not versions:
            continue
        ver_str = ', '.join(sorted(versions))
        for ii, dd in enumerate([pkg, ver_str]):
            if len(dd) > widths[ii]:
                widths[ii] = len(dd)
        lines.append([pkg, ver_str])

    if lines:
        build_table(widths, lines, headers)
    else:
        log.warning('No recognized package versions found.')


def print_env_vars():
    """Print the environment variables table."""
    headers = ['Environment Variable', 'Value']
    widths = np.array([len(h) for h in headers])
    lines = []
    for var in _ENV_VARS:
        val = os.environ.get(var, '')
        if not val:
            continue
        for ii, dd in enumerate([var, val]):
            if len(dd) > widths[ii]:
                widths[ii] = len(dd)
        lines.append([var, val])

    if lines:
        build_table(widths, lines, headers)
    else:
        log.warning('No expected environment variables are set.')


def main():

    p = argparse.ArgumentParser()
    p.add_argument('--specprod', type=str, required=True, help='spectroscopic production name (e.g., iron)')
    p.add_argument('--vacdir', type=str, required=True, help='top-level VAC directory')
    p.add_argument('--all-versions', action='store_true',
                   help='read every healpix file for code versions (slow; default: one per survey/program)')
    args = p.parse_args()

    all_files = {}
    for script in ('fastspec', 'fastphot'):
        files = gather_catalogs(args.vacdir, args.specprod, script)
        if len(files) > 0:
            all_files[script] = files

    if not all_files:
        log.error('No fastspec or fastphot catalogs found in {}/'.format(
            os.path.join(args.vacdir, 'catalogs')))
        return 1

    for script, files in all_files.items():
        print(f'\n### File Statistics ({script}) ###')
        print_file_statistics(files, args.specprod, script)

    for script, files in all_files.items():
        print(f'\n### QSO Redshift Corrections ({script}) ###')
        print_qso_statistics(files, args.specprod, script)

    print('\n### Software Versions ###\n')
    print_code_versions(args.vacdir, all_files, sample=not args.all_versions)

    print('\n### Environment Variables ###\n')
    print_env_vars()


if __name__ == '__main__':
    main()
