# reid_gallery_example.py
# Enrol vessels in a gallery, rank a query by cosine similarity and call it unknown below a threshold (NumPy, synthetic embeddings).
#
# By Sunil Tyagi (https://suniltyagi.in), for the note "Re-identifying vessels from radiated noise":
# https://suniltyagi.in/notes/vessel-re-identification/
#
# Nothing here is a trained model or real data. A real encoder maps each 5-second segment of a recording to a
# 512-number embedding; here, made-up vectors stand in for those embeddings, so that the gallery logic can be
# run and read on its own. The scores it prints come from this synthetic model, not from any measurement.
#
# Run: python reid_gallery_example.py   (Python 3 and NumPy only)

import numpy as np

DIM = 512           # embedding size
THRESHOLD = 0.50    # decision threshold for THIS synthetic example only; a real one is calibrated per deployment
rng = np.random.default_rng(7)


def unit(v):
    """L2-normalise a vector (or each row of a matrix) to length 1."""
    return v / np.linalg.norm(v, axis=-1, keepdims=True)


# --- A synthetic world: what the encoder would output -------------------------------------------------------
# Every vessel has its own direction (its identity). All ship noise shares a common component, so that two
# different vessels are not completely unlike each other. Each recording (passage) adds its own offset for the
# conditions it was made in, and each 5-second segment adds a little noise of its own.
common = unit(rng.normal(size=DIM))


def new_vessel():
    return unit(rng.normal(size=DIM))


def segment_embeddings(identity, n_segments=12, condition=0.55, noise=0.35):
    """The embeddings an encoder would give for the segments of one passage of a vessel."""
    offset = condition * unit(rng.normal(size=DIM))                 # this passage's recording conditions
    segs = identity + 0.6 * common + offset + noise * rng.normal(size=(n_segments, DIM)) / np.sqrt(DIM) * 8
    return unit(segs)                                                # the encoder outputs unit vectors


# --- The gallery -------------------------------------------------------------------------------------------
def passage_embedding(segs):
    """One recording's fingerprint: the mean of its segment embeddings, renormalised to length 1."""
    return unit(segs.mean(axis=0))


gallery_vecs, gallery_names = [], []


def enrol(name, passages):
    """Add a vessel's passages to the gallery. The model is not touched: enrolment only appends."""
    for segs in passages:
        gallery_vecs.append(passage_embedding(segs))
        gallery_names.append(name)


def identify(query_segs, top=3):
    """Score a query against every enrolled passage, keep each vessel's best passage, rank the vessels."""
    q = passage_embedding(query_segs)
    scores = np.array(gallery_vecs) @ q                              # cosine similarity: all vectors have length 1
    best = {}
    for name, s in zip(gallery_names, scores):
        best[name] = max(best.get(name, -1.0), float(s))
    ranked = sorted(best.items(), key=lambda kv: -kv[1])[:top]
    decision = ranked[0][0] if ranked[0][1] >= THRESHOLD else "unknown"
    return ranked, decision


def show(title, ranked, decision):
    print(title)
    for i, (name, s) in enumerate(ranked, 1):
        print(f"  {i}. {name:<8} cosine {s:.3f}")
    print(f"  decision: {decision} (threshold {THRESHOLD:.2f})\n")


vessels = {name: new_vessel() for name in ["ALPHA", "BRAVO", "CHARLIE", "DELTA", "ECHO"]}

# Enrol four vessels: several passages for some, one for others.
enrol("ALPHA", [segment_embeddings(vessels["ALPHA"]) for _ in range(3)])
enrol("BRAVO", [segment_embeddings(vessels["BRAVO"]) for _ in range(2)])
enrol("CHARLIE", [segment_embeddings(vessels["CHARLIE"]) for _ in range(1)])
enrol("DELTA", [segment_embeddings(vessels["DELTA"]) for _ in range(1)])
print(f"Gallery: {len(gallery_vecs)} passages of {len(set(gallery_names))} vessels, "
      f"each a {gallery_vecs[0].shape[0]}-number unit vector\n")

# 1. A new passage of an enrolled vessel.
show("Query 1: a new passage of BRAVO", *identify(segment_embeddings(vessels["BRAVO"])))

# 2. A vessel that was never enrolled.
echo_query = segment_embeddings(vessels["ECHO"])
show("Query 2: ECHO, never enrolled", *identify(echo_query))

# 3. Enrol ECHO from two passages. No retraining: the gallery grows by two vectors.
enrol("ECHO", [segment_embeddings(vessels["ECHO"]) for _ in range(2)])
print(f"Enrolled ECHO: the gallery now holds {len(gallery_vecs)} passages of {len(set(gallery_names))} vessels\n")
show("Query 3: the same ECHO passage as query 2", *identify(echo_query))

# 4. Why several passages per vessel help: the same query against one enrolled passage, then against three.
query = segment_embeddings(vessels["ALPHA"], condition=0.9)          # recorded in less familiar conditions
alpha_passages = [v for v, n in zip(gallery_vecs, gallery_names) if n == "ALPHA"]
q = passage_embedding(query)
one = float(alpha_passages[0] @ q)
three = max(float(v @ q) for v in alpha_passages)
print("Query 4: a passage of ALPHA in less familiar conditions")
for label, s in (("1 enrolled ALPHA passage ", one), ("3 enrolled ALPHA passages", three)):
    print(f"  against {label}: best cosine {s:.3f} -> {'ALPHA' if s >= THRESHOLD else 'unknown'}")