"""Fixed color-bin semantics and reproducible dataset membership."""
import hashlib
import io
import json
from pathlib import Path
import numpy as np
from PIL import Image
import pyarrow.parquet as pq
from scipy.spatial import cKDTree
from skimage.color import rgb2lab
import torch
from torch.utils.data import Dataset

class ColorBins:
    def __init__(self, centers, weights=None):
        self.centers = np.asarray(centers, np.float32)
        if self.centers.ndim != 2 or self.centers.shape[1] != 2:
            raise ValueError("Expected Q by 2 color centers")
        self.tree = cKDTree(self.centers)
        self.weights = np.ones(len(self.centers), np.float32) if weights is None else np.asarray(weights,np.float32)

    def encode(self, ab, k=5, sigma=5):
        k = min(k, len(self.centers))
        if sigma <= 0 or k < 1:
            raise ValueError("Invalid encoding parameters")
        dist, idx = self.tree.query(ab.reshape(-1,2), k=k)
        dist, idx = dist.reshape(-1,k), idx.reshape(-1,k)
        logw = -dist**2/(2*sigma**2)
        logw -= logw.max(1, keepdims=True)
        weight = np.exp(logw); weight /= weight.sum(1,keepdims=True)
        shape = (*ab.shape[:-1], k)
        return idx.reshape(shape).astype(np.int64), weight.reshape(shape).astype(np.float32)

def load_table(path):
    table = pq.read_table(path)
    return table

def validate_manifest(path, table, manifest):
    with open(path,'rb') as f:
        digest=hashlib.file_digest(f,'sha256').hexdigest()
    if digest!=manifest['dataset_sha256'] or len(table)!=manifest['rows']:
        raise ValueError('Dataset differs from pinned split manifest')
    groups=[list(map(int,manifest[k])) for k in ['train','validation','test']]
    if any(not g or len(g)!=len(set(g)) or min(g)<0 or max(g)>=len(table) for g in groups):
        raise ValueError('Invalid or repeated indices in manifest')
    if any(set(groups[i])&set(groups[j]) for i,j in [(0,1),(0,2),(1,2)]):
        raise ValueError('Train/validation/test overlap')

def decode_image(table, index):
    obj = table['image'][int(index)].as_py()
    return Image.open(io.BytesIO(obj['bytes'])).convert('RGB')

def legacy_split(n, seed=0):
    # Exactly reproduce datasets.Dataset.train_test_split(test_size=.05, seed=0).
    order = np.random.default_rng(seed).permutation(n)
    nval = int(np.ceil(.05*n))
    return order[nval:], order[:nval]

def stratified_subset(indices, labels, per_class, seed):
    rng = np.random.default_rng(seed)
    selected = []
    for label in sorted(set(labels)):
        group = np.asarray([i for i in indices if labels[i] == label])
        selected.extend(rng.choice(group, min(per_class,len(group)), replace=False).tolist())
    return np.asarray(selected, dtype=np.int64)

class PhotoDataset(Dataset):
    def __init__(self, table, indices, bins, size=256, augment=False):
        self.table, self.indices, self.bins = table, list(map(int,indices)), bins
        self.size, self.augment = size, augment

    def __len__(self):
        return len(self.indices)

    def __getitem__(self, item):
        # Keep square preprocessing for controlled comparison with historical runs.
        # Production inference preserves aspect ratio separately.
        image = decode_image(self.table,self.indices[item]).resize((self.size,self.size),Image.Resampling.BILINEAR)
        arr = np.asarray(image,dtype=np.float32)/255
        if self.augment and torch.rand(()) < .5:
            arr = arr[:,::-1].copy()
        lab = rgb2lab(arr).astype(np.float32)
        idx, weight = self.bins.encode(lab[...,1:])
        return (torch.from_numpy((lab[...,:1]/50-1).transpose(2,0,1).copy()),
                torch.from_numpy(lab[...,1:].transpose(2,0,1).copy()),
                torch.from_numpy(idx),torch.from_numpy(weight))

def estimate_weights(table, indices, bins, mix=.7, limit=1000, seed=0):
    """Re-estimate the prior ON checkpoint bins; never replace the bin vocabulary."""
    if not 0 <= mix <= 1:
        raise ValueError("Rebalance mixture must be in [0,1]")
    rng=np.random.default_rng(seed)
    chosen=rng.choice(indices,min(limit,len(indices)),replace=False)
    counts=np.ones(len(bins.centers),np.float64)*1e-3
    for index in chosen:
        image=decode_image(table,index).resize((32,32))
        ab=rgb2lab(np.asarray(image,dtype=np.float32)/255)[...,1:]
        # Low chroma is a heuristic, not proof that an image was originally B&W.
        if np.linalg.norm(ab,axis=-1).mean()<3:
            continue
        idx,weight=bins.encode(ab)
        np.add.at(counts,idx.ravel(),weight.ravel())
    prior=counts/counts.sum()
    weights=1/((1-mix)*prior+mix/len(prior))
    weights/=np.sum(prior*weights)
    return weights.astype(np.float32),prior.astype(np.float32)

def write_manifest(path, table, train, validation, test, dataset_sha256):
    groups=[set(map(int,g)) for g in [train,validation,test]]
    assert not groups[0]&groups[1] and not groups[0]&groups[2] and not groups[1]&groups[2]
    content={'dataset':'johnowhitaker/imagenette2-320','dataset_sha256':dataset_sha256,
        'rows':len(table),'split_seed':0,'legacy_holdout_fraction':.05,
        'train':list(map(int,train)),'validation':list(map(int,validation)), 'test':list(map(int,test)),
        'caveat':'Matches the recent seed-0 holdout; older upstream training exposure is not established.'}
    Path(path).write_text(json.dumps(content,indent=2))
    return content