#!/usr/bin/env python3
"""Check a shard set before spending a day training on it.

Cheap insurance.  The extraction is a four-hour pass over a filesystem that has failed
mid-run before, and every failure mode it has -- a truncated push, a day of download
stubs, a split that leaks -- is silent at training time and expensive to discover from a
loss curve.
"""

import argparse
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from suvi import dataset


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--shards", required=True)
    parser.add_argument("--manifest", default=None)
    parser.add_argument("--satellites", default="16,18")
    args = parser.parse_args(argv)

    satellites = tuple(int(s) for s in args.satellites.split(","))
    manifest_path = args.manifest or os.path.join(args.shards, "manifest.json")
    manifest = json.load(open(manifest_path))
    splits = manifest["splits"]

    print(f"manifest {manifest['digest']}  size={manifest['size']} "
          f"asinh={manifest.get('asinh_scale')}/{manifest.get('asinh_range')}")

    leaks = dataset.check_split(splits)
    print("split: " + ", ".join(f"{k} {len(v)}d" for k, v in sorted(splits.items())))
    print(f"  adjacent boundaries: {len(leaks)}{' ' + str(leaks[:3]) if leaks else ''}")

    problems = []
    totals = {}
    for name, days in sorted(splits.items()):
        slots = empty = missing = 0
        for day in days:
            for satellite in satellites:
                path = os.path.join(args.shards, dataset.shard_name(day, satellite))
                if not os.path.exists(path):
                    missing += 1
                    problems.append(f"missing shard {os.path.basename(path)}")
                    continue
                try:
                    with dataset.Shard(path) as shard:
                        times = shard.times()
                        usable = sum(1 for t in times if shard.raw(t) is not None)
                        slots += usable
                        if usable == 0:
                            empty += 1
                            problems.append(f"empty shard {os.path.basename(path)}")
                        # Spot-check one decode per shard.
                        for when in times[:1]:
                            frames = shard.coded(when)
                            if frames is None:
                                continue
                            if frames.shape != (6, dataset.SHARD_SIZE, dataset.SHARD_SIZE):
                                problems.append(f"{os.path.basename(path)}: {frames.shape}")
                            if abs(frames).max() > 1.0001:
                                problems.append(
                                    f"{os.path.basename(path)}: out of range "
                                    f"{abs(frames).max():.3f}")
                except (ValueError, OSError) as error:
                    problems.append(f"{os.path.basename(path)}: {error}")
        totals[name] = slots
        print(f"  {name:>5}: {slots:>6} usable slots"
              + (f", {empty} empty shards" if empty else "")
              + (f", {missing} MISSING" if missing else ""))

    print(f"\ntotal usable slots: {sum(totals.values())}")
    if problems:
        print(f"\n{len(problems)} problems:")
        for line in problems[:20]:
            print(f"  {line}")
        return 1
    print("no problems found")
    return 0


if __name__ == "__main__":
    sys.exit(main())