"""MOA fingerprint matching by cosine similarity."""

from __future__ import annotations

import numpy as np
import pandas as pd


FEATURE_COLS = [
    "macro_pct",
    "medium_pct",
    "micro_pct",
    "total_LD_area_per_cell",
    "ld_count_per_cell",
    "manders_m1",
    "plin1_mean",
    "plin5_mean",
]


def _zscale(reference: pd.DataFrame, query: pd.DataFrame, cols=FEATURE_COLS):
    """Z-score scale based on reference table; apply identically to query."""
    mu = reference[cols].mean()
    sd = reference[cols].std().replace(0, 1.0)
    ref_z = (reference[cols] - mu) / sd
    qry_z = (query[cols] - mu) / sd
    return ref_z, qry_z


def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    na = np.linalg.norm(a)
    nb = np.linalg.norm(b)
    if na == 0 or nb == 0:
        return float("nan")
    return float(np.dot(a, b) / (na * nb))


def match_moa(
    well_summary: pd.DataFrame,
    reference: pd.DataFrame,
    feature_cols=FEATURE_COLS,
) -> pd.DataFrame:
    """For every well row in `well_summary`, score similarity to each reference drug.

    Returns long table: well, drug(observed), dose_uM, ref_drug, similarity, top1.
    """
    cols = [c for c in feature_cols if c in well_summary.columns and c in reference.columns]
    if not cols:
        return pd.DataFrame()

    ref_z, qry_z = _zscale(reference, well_summary, cols)

    rows = []
    for q_idx, q in qry_z.iterrows():
        q_vec = q.values.astype(float)
        meta = well_summary.iloc[q_idx]
        for r_idx, r in ref_z.iterrows():
            r_vec = r.values.astype(float)
            sim = _cosine(q_vec, r_vec)
            ref_meta = reference.iloc[r_idx]
            rows.append(
                {
                    "well": meta.get("well", ""),
                    "drug_observed": meta.get("drug", ""),
                    "dose_uM": meta.get("dose_uM", float("nan")),
                    "ref_drug": ref_meta.get("drug", ""),
                    "similarity": sim,
                }
            )
    out = pd.DataFrame(rows)
    if out.empty:
        return out
    # Mark top-1 reference for each well
    out["rank"] = out.groupby("well")["similarity"].rank(ascending=False, method="first")
    out["top1"] = out["rank"] == 1
    return out