"""Fresh full-patch inference for the October 12 frozen spatial-context methods."""
import argparse, json, time
from pathlib import Path
import numpy as np
import torch
import config as C
import foundation_transfer as F
from dataset import PASTISPatchDataset, pad_collate, val_patch_ids, test_patch_ids, crop_ij
from notebook_ensemble import CompleteEnsemble
from spatial_context_transfer import SpatialContextTransfer

class CurrentEnsemble(CompleteEnsemble):
    def __init__(self, paths):
        paths = dict(paths)
        folder = paths.pop('spatial_context')
        super().__init__(paths)
        self.context = SpatialContextTransfer(folder)

    def predict(self, x, pos, days, split):
        parts = super().predict(x, pos, days, split, return_features=True)
        return self.context.predict(parts['final'], parts['optical_features'], x.shape[-1])

def evaluate(paths, split='val', output=None, max_batches=0, save_probabilities=False):
    """A limited smoke run is explicitly marked incomplete and has no headline score."""
    from torch.utils.data import DataLoader
    assert split in ['val', 'test']
    torch.set_num_threads(4); torch.manual_seed(3407)
    torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True
    ids = val_patch_ids() if split == 'val' else test_patch_ids()
    dl = DataLoader(PASTISPatchDataset(ids), batch_size=2 if split == 'val' else 1,
                    shuffle=False, num_workers=2, pin_memory=True, collate_fn=pad_collate)
    model = CurrentEnsemble(paths); matrices = {}; offset = 0; entries = 0; chunks = []
    started = time.monotonic()
    try:
        for bi, ((x, pos, days), y) in enumerate(dl):
            if max_batches and bi >= max_batches: break
            x, pos = x.cuda(), pos.cuda(); entries += len(x)
            for i, j in C.FT_CROP_IJ if split == 'val' else [(0, 0)]:
                crop, target = crop_ij(x, y, i, j, 4 if split == 'val' else 1)
                parts = model.predict(crop, pos, days, split)
                for name, p in parts.items():
                    matrices.setdefault(name, np.zeros((20, 20), np.int64))
                    matrices[name] += F.confusion(target.numpy().ravel(), p.argmax(-1).ravel())
                if save_probabilities: chunks.append(parts['best_observed'].reshape(-1, 20))
                offset += target.numel()
            if bi % 10 == 0: print(split, bi, len(dl), 'seconds', time.monotonic() - started, flush=True)
    finally:
        model.close()
    complete = entries == len(ids)
    result = {'split': split, 'complete': complete, 'patch_entries': entries,
              'expected_patch_entries': len(ids), 'seed': 3407, 'pixel_entries': offset,
              'seconds': time.monotonic() - started,
              'components': {k: F.metrics(v) for k, v in matrices.items()},
              'paths': {k: str(v) for k, v in paths.items()},
              'selection': 'validation_selected is the primary result; best_observed is exploratory test-shortlist winner'}
    if output:
        output = Path(output); output.mkdir(parents=True, exist_ok=True)
        (output / f'fresh_{split}_results.json').write_text(json.dumps(result, indent=2))
        if save_probabilities: np.save(output / f'fresh_{split}_probabilities.npy', np.concatenate(chunks))
    return result


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--paths', type=Path, required=True)
    parser.add_argument('--split', choices=['val','test'], default='val')
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--max_batches', type=int, default=0)
    parser.add_argument('--save_probabilities', action='store_true')
    a=parser.parse_args()
    r=evaluate(json.loads(a.paths.read_text()),a.split,a.output,a.max_batches,a.save_probabilities)
    print(json.dumps(r,indent=2))