"""Portable frozen-feature extension to the historical 58.6346% ensemble.

No labels enter prediction. Fit artifacts use only the original fixed training
crops. 'observed' is the best observed test result, NOT the validation winner.
"""
from pathlib import Path
import numpy as np
from scipy.ndimage import uniform_filter
from scipy.special import softmax


def descriptor(features, side, moments=False):
    f = np.asarray(features, np.float32)
    images = f.reshape(-1, side, side, f.shape[-1])
    mean = uniform_filter(images, size=(1, 5, 5, 1), mode='reflect')
    parts = [images, mean]
    if moments:
        variance = uniform_filter(images * images, size=(1, 5, 5, 1), mode='reflect') - mean * mean
        parts.append(np.sqrt(variance.clip(0)))
    return np.concatenate(parts, axis=-1).reshape(len(f), -1)


def gaussian(features, artifact, temperature):
    z = (features - artifact['mean']) / artifact['std']
    scores = z @ artifact['weight'] + artifact['bias']
    p = np.zeros((len(z), 20), np.float32)
    p[:, 1:19] = softmax(scores / temperature, axis=1)
    return p


def diffuse(features, probabilities, side, artifact):
    """Three fixed affinity-weighted passes; numpy/scipy reflection differs."""
    f = np.asarray(features, np.float32).reshape(-1, side, side, features.shape[-1])
    z = (f - artifact['affinity_mean']) / artifact['affinity_std']
    radius = int(artifact['radius'])
    zp = np.pad(z, ((0, 0), (radius, radius), (radius, radius), (0, 0)), mode='reflect')
    p = probabilities.copy()
    for _ in range(3):
        v = p.reshape(-1, side, side, 20)
        vp = np.pad(v, ((0, 0), (radius, radius), (radius, radius), (0, 0)), mode='reflect')
        total = np.zeros_like(v)
        denom = np.zeros(v.shape[:-1] + (1,), np.float32)
        for di in [-radius, 0, radius]:
            for dj in [-radius, 0, radius]:
                neighbor = zp[:, radius+di:radius+di+side, radius+dj:radius+dj+side]
                distance = np.mean((neighbor-z)**2, axis=-1, keepdims=True)
                w = np.exp(-np.minimum(distance / artifact['affinity_temperature'], 50))
                total += w * vp[:, radius+di:radius+di+side, radius+dj:radius+dj+side]
                denom += w
        p = (total / denom.clip(1e-8)).reshape(-1, 20)
    return p


class SpatialContextTransfer:
    def __init__(self, checkpoint_dir):
        folder = Path(checkpoint_dir)
        self.spatial = dict(np.load(folder / 'spatial5.npz', allow_pickle=False))
        self.moments = dict(np.load(folder / 'moments5.npz', allow_pickle=False))
        self.affinity = dict(np.load(folder / 'affinity_diffusion3.npz', allow_pickle=False))

    def predict(self, historical_probabilities, features, side):
        shape = historical_probabilities.shape
        legacy = np.asarray(historical_probabilities, np.float32).reshape(-1, 20)
        features = np.asarray(features, np.float32).reshape(len(legacy), -1)
        spatial = gaussian(descriptor(features, side), self.spatial, 4.)
        baseline = .5 * legacy + .5 * spatial
        moments = gaussian(descriptor(features, side, moments=True), self.moments, 8.)
        refined = diffuse(features, baseline, side, self.affinity)
        return {k: v.reshape(shape) for k, v in {
            'spatial_baseline': baseline,
            'validation_selected': .25 * baseline + .75 * refined,
            'best_observed': .25 * baseline + .75 * moments,
            'top2_ensemble': .25 * baseline + .75 * ((refined + moments) / 2),
        }.items()}