"""Build the final training CSV by combining:
  * siRNA sequences + knockdown labels (from synthetic or real input)
  * Reynolds 2004 efficacy features (from `reynolds_features.py`)
  * Off-target risk scores against the 6 non-target species
    (using the `KmerOffTargetIndex` from `_legacy_imports.py`)

Output columns:
    sirna_seq, target_gene, knockdown_pct, pest_label,
    gc_content, no_repeats, at_pos19, a_pos3, t_pos10, ag_pos13,
    t_pos16, thermo_asymmetry, reynolds_score,
    offtarget_apis_mellifera, offtarget_bos_taurus,
    offtarget_bos_indicus, offtarget_gallus_gallus,
    offtarget_danio_rerio, offtarget_homo_sapiens

`pest_label` is the binarized knockdown_pct (1 if KD >= 0.7 else 0).
This is the SCIENTIFICALLY CORRECT label (the original merged
biopesticide AI code wrongly labeled any 100-nt tile from a pest gene
as pest_label=1, which has no basis in efficacy).

Usage (run from the project root):

    # Synthetic path (works without any external data):
    python -m scripts.data.build_training_csv --source synthetic

    # Real path (requires real siRNA CSV + NCBI rna.fna files on disk):
    python -m scripts.data.build_training_csv --source real
    python -m scripts.data.build_training_csv --source real \
        --input /path/to/your/sirna_real.csv

Dependencies: pandas, numpy (transitively via reynolds_features), tqdm.
"""

from __future__ import annotations

import argparse
import logging
import sys
from collections import defaultdict
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

import pandas as pd

try:
    from tqdm import tqdm
except ImportError:  # pragma: no cover
    def tqdm(iterable, **_kwargs):
        return iterable

# Legacy helpers (copied from upload/merged_biopesticide_ai.py to avoid
# depending on the soon-to-be-refactored `src/` package layout).
from scripts.data._legacy_imports import (
    KmerOffTargetIndex,
    SAFETY_SPECIES,
    TARGET_SPECIES,
    fasta_iter,
    generate_kmers,
)
from scripts.data.reynolds_features import ReynoldsFeaturizer


# --------------------------------------------------------------------------- #
# Logging + paths
# --------------------------------------------------------------------------- #
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("build_training_csv")

PROJECT_ROOT = Path(__file__).resolve().parents[2]

# Canonical 6-species safety panel (same order as the output CSV columns).
SAFETY_PANEL: Tuple[str, ...] = tuple(SAFETY_SPECIES)

# Default I/O paths.
SYNTHETIC_INPUT = PROJECT_ROOT / "data" / "synthetic" / "sirna_training.csv"
SYNTHETIC_SAFETY_FASTA = PROJECT_ROOT / "data" / "synthetic" / "safety_transcripts.fasta"
SYNTHETIC_OUTPUT = PROJECT_ROOT / "data" / "processed" / "training_data.csv"

REAL_INPUT_DEFAULT = PROJECT_ROOT / "data" / "external" / "sirna_real.csv"
REAL_OUTPUT = PROJECT_ROOT / "data" / "processed" / "training_data_real.csv"

# Output column order (spec-defined, dynamic per safety species count).
_BASE_COLUMNS: Tuple[str, ...] = (
    "sirna_seq",
    "target_gene",
    "knockdown_pct",
    "pest_label",
    "gc_content",
    "no_repeats",
    "at_pos19",
    "a_pos3",
    "t_pos10",
    "ag_pos13",
    "t_pos16",
    "thermo_asymmetry",
    "reynolds_score",
)
OUTPUT_COLUMNS: Tuple[str, ...] = _BASE_COLUMNS + tuple(
    f"offtarget_{sp}" for sp in SAFETY_SPECIES
)

# K-mer size used by the off-target index. Must match the legacy default.
KMER_K = 21

# Binaraization threshold (task spec).
PEST_LABEL_THRESHOLD = 0.7

# NCBI species -> assembly id, for the real-data safety panel. Matches
# `download_sources.py`. The pest species (nilaparvata_lugens) is excluded
# from the safety panel.
NCBI_SPECIES_ASSEMBLIES: Dict[str, str] = {
    "apis_mellifera": "GCF_003254395.2",
    "bos_taurus": "GCF_002263795.1",
    "bos_indicus": "GCA_014661045.1",
    "gallus_gallus": "GCF_016699485.2",
    "danio_rerio": "GCF_000002035.4",
    "homo_sapiens": "GCF_000001405.40",
}


# --------------------------------------------------------------------------- #
# Indexing helpers
# --------------------------------------------------------------------------- #
def index_synthetic_safety_panel(
    idx: KmerOffTargetIndex, fasta_path: Path
) -> Dict[str, int]:
    """Build the off-target index from the synthetic safety FASTA.

    The synthetic FASTA groups 5 species (skip human) into a single
    file with headers like `>apis_mellifera_fake_001`. We group records
    by their species prefix (everything before `_fake_`) and feed each
    group into the index.

    Returns a dict {species_name: number_of_records_indexed}.
    """
    by_species: Dict[str, List[str]] = defaultdict(list)
    for header, seq in fasta_iter(fasta_path):
        # header is e.g. 'apis_mellifera_fake_001' -> species 'apis_mellifera'
        if "_fake_" in header:
            species = header.split("_fake_", 1)[0]
        else:
            # Fallback: take the first two underscore-separated tokens.
            parts = header.split("_")
            species = "_".join(parts[:2]) if len(parts) >= 2 else header
        by_species[species].append(seq)

    counts: Dict[str, int] = {}
    for species, seqs in by_species.items():
        species_set = set()
        for seq in seqs:
            for kmer in generate_kmers(seq, idx.k):
                idx.index[kmer] += 1
                species_set.add(kmer)
        idx.species_kmers[species] = species_set
        counts[species] = len(seqs)
        log.info(
            "  Indexed %-20s %4d transcripts -> %d unique %d-mers",
            species,
            len(seqs),
            len(species_set),
            idx.k,
        )
    return counts


def index_real_safety_panel(idx: KmerOffTargetIndex) -> Dict[str, Path]:
    """Build the off-target index from on-disk NCBI rna.fna files.

    Missing files are logged and skipped; the corresponding species'
    off-target column will be 0.0 in the output.
    """
    found: Dict[str, Path] = {}
    for species, assembly in NCBI_SPECIES_ASSEMBLIES.items():
        rna_fna = (
            PROJECT_ROOT
            / "data"
            / species
            / "ncbi_dataset"
            / "data"
            / assembly
            / "rna.fna"
        )
        if not rna_fna.is_file():
            log.warning(
                "  Missing NCBI file for %s: %s -> offtarget_%s will be 0.0",
                species,
                rna_fna,
                species,
            )
            continue
        log.info("  Indexing %s from %s", species, rna_fna)
        idx.build_from_fasta(rna_fna, species)
        found[species] = rna_fna
    return found


# --------------------------------------------------------------------------- #
# Input loading + validation
# --------------------------------------------------------------------------- #
def load_input_siRNAs(path: Path) -> pd.DataFrame:
    """Load the input siRNA CSV. Required columns: sirna_seq, target_gene, knockdown_pct.

    Extra columns (e.g. `source`) are kept but not used.
    """
    if not path.is_file():
        raise FileNotFoundError(
            f"Input siRNA CSV not found: {path}\n"
            f"  For --source synthetic, run `python -m scripts.data.generate_synthetic` first.\n"
            f"  For --source real, place a CSV at {REAL_INPUT_DEFAULT} or pass --input <path>."
        )
    df = pd.read_csv(path)
    required = {"sirna_seq", "target_gene", "knockdown_pct"}
    missing = required - set(df.columns)
    if missing:
        raise ValueError(
            f"Input CSV {path} is missing required columns: {sorted(missing)}. "
            f"Found columns: {list(df.columns)}"
        )
    # Coerce knockdown_pct to float; validate range.
    df["knockdown_pct"] = pd.to_numeric(df["knockdown_pct"], errors="coerce")
    if df["knockdown_pct"].isna().any():
        raise ValueError(f"Input CSV {path} has non-numeric knockdown_pct values.")
    if ((df["knockdown_pct"] < 0.0) | (df["knockdown_pct"] > 1.0)).any():
        log.warning("  Some knockdown_pct values are outside [0, 1]; clipping.")
        df["knockdown_pct"] = df["knockdown_pct"].clip(0.0, 1.0)
    # Strip whitespace from sequence + uppercase + T-normalize (DNA).
    df["sirna_seq"] = (
        df["sirna_seq"].astype(str).str.strip().str.upper().str.replace("U", "T")
    )
    # Validate length.
    bad_lens = df[df["sirna_seq"].str.len() != 21]
    if not bad_lens.empty:
        raise ValueError(
            f"Input CSV has {len(bad_lens)} siRNAs that are not 21 nt long. "
            f"First few: {bad_lens['sirna_seq'].head(5).tolist()}"
        )
    # Drop exact duplicate sequences (keep first).
    n_before = len(df)
    df = df.drop_duplicates(subset=["sirna_seq"], keep="first").reset_index(drop=True)
    if len(df) < n_before:
        log.info("  Dropped %d duplicate siRNA sequences.", n_before - len(df))
    return df


# --------------------------------------------------------------------------- #
# Main build
# --------------------------------------------------------------------------- #
def build(
    source: str,
    input_csv: Path,
    output_csv: Path,
) -> None:
    log.info("=" * 70)
    log.info("Build configuration:")
    log.info("  source:       %s", source)
    log.info("  input CSV:    %s", input_csv)
    log.info("  output CSV:   %s", output_csv)
    log.info("  K-mer size:   %d", KMER_K)
    log.info("  KD threshold: %.2f (>= -> pest_label=1)", PEST_LABEL_THRESHOLD)
    log.info("=" * 70)

    # ---- 1. Load input siRNAs ------------------------------------------ #
    log.info("STEP 1/4: Loading input siRNAs.")
    df_in = load_input_siRNAs(input_csv)
    log.info("  Loaded %d unique siRNAs.", len(df_in))

    # ---- 2. Build off-target index ------------------------------------- #
    log.info("STEP 2/4: Building off-target index over safety panel.")
    idx = KmerOffTargetIndex(k=KMER_K)
    if source == "synthetic":
        if not SYNTHETIC_SAFETY_FASTA.is_file():
            raise FileNotFoundError(
                f"Synthetic safety FASTA not found: {SYNTHETIC_SAFETY_FASTA}\n"
                f"  Run `python -m scripts.data.generate_synthetic` first."
            )
        log.info("  Indexing synthetic safety panel: %s", SYNTHETIC_SAFETY_FASTA)
        counts = index_synthetic_safety_panel(idx, SYNTHETIC_SAFETY_FASTA)
        missing = [sp for sp in SAFETY_PANEL if sp not in idx.species_kmers]
        if missing:
            log.info(
                "  Synthetic path intentionally skips: %s "
                "(these columns will be 0.0 in the output).",
                ", ".join(missing),
            )
    else:  # real
        log.info("  Indexing real NCBI safety panel.")
        found = index_real_safety_panel(idx)
        missing = [sp for sp in SAFETY_PANEL if sp not in idx.species_kmers]
        if missing:
            log.warning(
                "  Real path is missing NCBI data for: %s "
                "(these columns will be 0.0 in the output).",
                ", ".join(missing),
            )
        if not found:
            log.warning(
                "  No NCBI safety FASTAs found on disk. Proceeding with all "
                "off-target columns set to 0.0. Run "
                "`python -m scripts.data.download_sources` to verify paths, "
                "or place real rna.fna files under data/<species>/ncbi_dataset/."
            )

    # ---- 3. Featurize + off-target scoring ----------------------------- #
    log.info("STEP 3/4: Featurizing siRNAs (Reynolds + off-target).")
    featurizer = ReynoldsFeaturizer()
    rows: List[dict] = []
    for _, row in tqdm(df_in.iterrows(), total=len(df_in), desc="siRNAs"):
        seq = row["sirna_seq"]
        kd = float(row["knockdown_pct"])
        feats = featurizer.featurize(seq)
        risks = idx.per_species_risk(seq)
        out_row = {
            "sirna_seq": seq,
            "target_gene": row["target_gene"],
            "knockdown_pct": kd,
            "pest_label": int(kd >= PEST_LABEL_THRESHOLD),
        }
        out_row.update(feats)
        for sp in SAFETY_PANEL:
            out_row[f"offtarget_{sp}"] = float(risks.get(sp, 0.0))
        rows.append(out_row)

    df_out = pd.DataFrame(rows, columns=list(OUTPUT_COLUMNS))

    # ---- 4. Write output ----------------------------------------------- #
    log.info("STEP 4/4: Writing output CSV.")
    output_csv.parent.mkdir(parents=True, exist_ok=True)
    df_out.to_csv(output_csv, index=False)
    log.info("  Wrote %d rows to %s", len(df_out), output_csv)

    # ---- Summary ------------------------------------------------------- #
    log.info("-" * 70)
    log.info("Build complete.")
    log.info("  Total siRNAs:      %d", len(df_out))
    log.info("  pest_label=1:      %d (%.1f%%)",
             int(df_out["pest_label"].sum()),
             100.0 * df_out["pest_label"].mean())
    log.info("  pest_label=0:      %d (%.1f%%)",
             int((1 - df_out["pest_label"]).sum()),
             100.0 * (1 - df_out["pest_label"]).mean())
    log.info("  Mean Reynolds score: %.2f / 8", df_out["reynolds_score"].mean())
    log.info("  Off-target column non-zero counts:")
    for sp in SAFETY_PANEL:
        col = f"offtarget_{sp}"
        nz = int((df_out[col] > 0).sum())
        log.info("    %-28s %4d / %d non-zero", col, nz, len(df_out))


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #
def parse_args(argv: Iterable[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="python -m scripts.data.build_training_csv",
        description=(
            "Build the final training CSV by combining siRNA efficacy "
            "labels, Reynolds 2004 features, and off-target risk against "
            "the 6-species safety panel."
        ),
    )
    p.add_argument(
        "--source",
        choices=("synthetic", "real"),
        required=True,
        help="Data source: 'synthetic' uses data/synthetic/, 'real' uses "
             "data/external/sirna_real.csv + on-disk NCBI rna.fna files.",
    )
    p.add_argument(
        "--input",
        type=Path,
        default=None,
        help="Override the input siRNA CSV path. Defaults to "
             "data/synthetic/sirna_training.csv (synthetic) or "
             "data/external/sirna_real.csv (real).",
    )
    p.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Override the output CSV path. Defaults to "
             "data/processed/training_data.csv (synthetic) or "
             "data/processed/training_data_real.csv (real).",
    )
    return p.parse_args(argv)


def main() -> int:
    args = parse_args()
    if args.input is not None:
        input_csv = args.input
    else:
        input_csv = SYNTHETIC_INPUT if args.source == "synthetic" else REAL_INPUT_DEFAULT
    if args.output is not None:
        output_csv = args.output
    else:
        output_csv = SYNTHETIC_OUTPUT if args.source == "synthetic" else REAL_OUTPUT
    try:
        build(source=args.source, input_csv=input_csv, output_csv=output_csv)
    except FileNotFoundError as exc:
        log.error("%s", exc)
        return 2
    except (ValueError, RuntimeError) as exc:
        log.error("Build failed: %s", exc)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())