# Log12 予備実験：PatchCore（WideResNet50）で配膳ロボットの長い止まりを見つけられるか
# 覚えさせるのは「普段」（往復の途中で走っている所）の切り出しだけ。
# 確かめは、別の回の 普段・短い止まり（3〜30秒）・長い止まり（30秒以上、始まりから15秒後）。
# 使い方: python3 patchcore.py <crops のフォルダ> <A か B> <結果のフォルダ>
import sys, os, glob, json, csv, random
import numpy as np, torch, torch.nn.functional as F
from torchvision.models import wide_resnet50_2, Wide_ResNet50_2_Weights
from PIL import Image
from sklearn.metrics import roc_auc_score

root, kind, out = sys.argv[1], sys.argv[2], sys.argv[3]
os.makedirs(out, exist_ok=True)
torch.manual_seed(0); random.seed(0); np.random.seed(0)
torch.set_num_threads(os.cpu_count())

model = wide_resnet50_2(weights=Wide_ResNet50_2_Weights.IMAGENET1K_V1).eval()
feats = {}
model.layer2.register_forward_hook(lambda m, i, o: feats.__setitem__('l2', o))
model.layer3.register_forward_hook(lambda m, i, o: feats.__setitem__('l3', o))
MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)

def load(paths):
    x = np.stack([np.asarray(Image.open(p).convert('RGB'), dtype=np.float32) / 255 for p in paths])
    return (torch.from_numpy(x).permute(0, 3, 1, 2) - MEAN) / STD

@torch.no_grad()
def embed(paths, bs=16):
    """1枚を 28x28 の区画に分け、区画ごとの特徴（layer2+layer3、3x3でならす）を返す"""
    outs = []
    for i in range(0, len(paths), bs):
        model(load(paths[i:i + bs]))
        l2 = F.avg_pool2d(feats['l2'], 3, 1, 1)
        l3 = F.interpolate(F.avg_pool2d(feats['l3'], 3, 1, 1), size=l2.shape[-2:], mode='bilinear', align_corners=False)
        outs.append(torch.cat([l2, l3], 1))   # (b, 1536, 28, 28)
    return torch.cat(outs)

def coreset(X, n, proj_dim=128):
    """貪欲法で代表の区画だけ残す（PatchCore の記憶を小さくする）"""
    P = X @ torch.randn(X.shape[1], proj_dim) / proj_dim ** 0.5
    idx = [random.randrange(len(X))]
    d = torch.cdist(P, P[idx]).squeeze(1)
    for _ in range(n - 1):
        j = int(torch.argmax(d)); idx.append(j)
        d = torch.minimum(d, torch.cdist(P, P[j:j + 1]).squeeze(1))
    return X[idx]

def score(paths, bank, bs=8):
    """区画ごとに、記憶の中で一番近い区画までの距離。1枚の点数は最大値。地図も返す"""
    sc, maps = [], []
    for i in range(0, len(paths), bs):
        E = embed(paths[i:i + bs]); b, c, h, w = E.shape
        Q = E.permute(0, 2, 3, 1).reshape(-1, c)
        dmin = torch.cat([torch.cdist(Q[k:k + 4096], bank).min(1).values for k in range(0, len(Q), 4096)])
        m = dmin.view(b, h, w)
        sc += m.flatten(1).max(1).values.tolist(); maps.append(m)
    return np.array(sc), torch.cat(maps).numpy()

d = f'{root}/{kind}'
train = sorted(glob.glob(f'{d}/train/normal/*.png'))
test = {lab: sorted(glob.glob(f'{d}/test/{lab}/*.png')) for lab in ('normal', 'short', 'long')}
print(kind, 'train', len(train), {k: len(v) for k, v in test.items()}, flush=True)

# 区画を全部持つとメモリが足りないので、1枚あたり200区画を無作為に残してから代表を選ぶ
Xs = []
for i in range(0, len(train), 16):
    E = embed(train[i:i + 16]); c = E.shape[1]
    Q = E.permute(0, 2, 3, 1).reshape(E.shape[0], -1, c)
    for q in Q: Xs.append(q[torch.randperm(q.shape[0])[:200]])
X = torch.cat(Xs); n_all = len(train) * 28 * 28
bank = coreset(X, int(n_all * 0.01))
print('bank', tuple(bank.shape), flush=True)

res, maps = {}, {}
for lab, ps in test.items():
    res[lab], maps[lab] = score(ps, bank)

def auc(a, b):
    return float(roc_auc_score([0] * len(a) + [1] * len(b), np.r_[a, b]))
# しきい値は 普段 の確かめ画像の 95% 点（普段のうち5%を誤報にする置き方）
thr = float(np.quantile(res['normal'], 0.95))
summary = {
    'kind': kind, 'n_train': len(train), 'n_test': {k: len(v) for k, v in test.items()}, 'bank': list(bank.shape),
    'auroc_long_vs_normal': auc(res['normal'], res['long']),
    'auroc_short_vs_normal': auc(res['normal'], res['short']),
    'auroc_long_vs_short': auc(res['short'], res['long']),
    'threshold_p95_normal': thr,
    'hit_rate_long_at_thr': float((res['long'] > thr).mean()),
    'hit_rate_short_at_thr': float((res['short'] > thr).mean()),
    'false_alarm_normal_at_thr': float((res['normal'] > thr).mean()),
    'score_median': {k: float(np.median(v)) for k, v in res.items()},
}
print(json.dumps(summary, ensure_ascii=False, indent=1), flush=True)
json.dump(summary, open(f'{out}/summary_{kind}.json', 'w'), ensure_ascii=False, indent=1)
with open(f'{out}/scores_{kind}.csv', 'w', newline='') as fo:
    w = csv.writer(fo); w.writerow(['label', 'file', 'score'])
    for lab, ps in test.items():
        for p, s in zip(ps, res[lab]): w.writerow([lab, os.path.basename(p), round(float(s), 4)])
np.savez_compressed(f'{out}/maps_{kind}.npz', **{lab: maps[lab] for lab in maps},
                    **{f'{lab}_files': np.array([os.path.basename(p) for p in ps]) for lab, ps in test.items()})