#!/usr/bin/env python3
"""
K-mer-based group prediction for unknown sequences.

Inputs:
- Unknown sequences: a FASTA file or a directory of FASTA files
- Unique k-mers: either
  * a directory containing unique_k{k}_{group}.tsv/.txt files (from script #1), OR
  * a ZIP file containing those files

Modes:
- fast: exact substring matching only (very fast)
- full: alignment-based matching (slower, more tolerant) + Fisher + FDR

Outputs:
- predictions_by_alignment.csv
- predicted_results_summary.png

Example:
  python kmer_predict.py \
    --unknown unknown_fastas/ \
    --kmer-input kmer_results.zip \
    --outdir pred_out \
    --seqtype dna \
    --mode fast

"""

from __future__ import annotations

import argparse
import os
import re
import shutil
import tempfile
import zipfile
from dataclasses import dataclass
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from scipy.stats import fisher_exact
from statsmodels.stats.multitest import multipletests

from Bio import Align
from Bio.Align import substitution_matrices


FASTA_EXTS = (".fasta", ".fa", ".fas", ".fna")
KMER_FILE_EXTS = (".tsv", ".txt")
DEFAULT_GROUP_REGEX = r"unique_k\d+_(.+)\.(tsv|txt)$"

BLOSUM62 = substitution_matrices.load("BLOSUM62")


# ----------------------------
# FASTA utilities
# ----------------------------

def read_fasta(filepath: str) -> Tuple[List[str], List[str]]:
    headers, seqs, seq = [], [], []
    with open(filepath, "r", encoding="utf-8") as fh:
        for line in fh:
            line = line.rstrip("\n")
            if not line:
                continue
            if line.startswith(">"):
                if seq:
                    seqs.append("".join(seq))
                    seq = []
                headers.append(line[1:].strip())
            else:
                seq.append(line.strip().upper())
        if seq:
            seqs.append("".join(seq))
    return headers, seqs


def clean_protein(seq: str) -> str:
    return re.sub(r"[^ACDEFGHIKLMNPQRSTVWY]", "", seq.upper())


def clean_dna(seq: str) -> str:
    # allow U and N like your original
    return re.sub(r"[^ACGTUN]", "", seq.upper())


def iter_unknown_sequences(unknown: str, is_protein: bool) -> List[Tuple[str, str, str]]:
    """
    Returns list of (source_file, header, cleaned_seq).
    unknown can be a fasta file or a directory with fasta files.
    """
    seq_index: List[Tuple[str, str, str]] = []

    if os.path.isdir(unknown):
        files = [
            os.path.join(unknown, f)
            for f in os.listdir(unknown)
            if f.lower().endswith(FASTA_EXTS)
        ]
    else:
        files = [unknown]

    files = [f for f in files if os.path.isfile(f)]
    for fp in sorted(files):
        headers, seqs = read_fasta(fp)
        if is_protein:
            seqs = [clean_protein(s) for s in seqs]
        else:
            seqs = [clean_dna(s) for s in seqs]

        for h, s in zip(headers, seqs):
            if s:  # drop empty after cleaning
                seq_index.append((fp, h, s))

    return seq_index


# ----------------------------
# ZIP utilities (safe extract)
# ----------------------------

def safe_extract_zip(zip_path: str, dst_dir: str) -> None:
    """Extract ZIP safely (prevents zip-slip)."""
    with zipfile.ZipFile(zip_path, "r") as z:
        for member in z.infolist():
            if member.is_dir():
                continue
            target = os.path.normpath(os.path.join(dst_dir, member.filename))
            if not target.startswith(os.path.abspath(dst_dir) + os.sep):
                continue  # skip suspicious paths
            os.makedirs(os.path.dirname(target), exist_ok=True)
            with z.open(member) as src, open(target, "wb") as out:
                shutil.copyfileobj(src, out)


# ----------------------------
# Load unique kmers
# ----------------------------

@dataclass
class KmerDB:
    group_kmers: Dict[str, List[str]]
    source_dir: str


def parse_group_from_filename(fname: str, group_regex: str) -> str:
    m = re.search(group_regex, fname, re.IGNORECASE)
    if m:
        return m.group(1)
    # fallback: remove extension
    return os.path.splitext(fname)[0]


def load_unique_kmers_from_dir(
    kmer_dir: str,
    is_protein: bool,
    group_regex: str = DEFAULT_GROUP_REGEX,
) -> KmerDB:
    """
    Loads k-mers from files like:
      unique_k15_group1.tsv
      unique_k20_groupA.txt
    Accepts TSV or TXT; ignores comment/header lines.
    """
    group_kmers: Dict[str, List[str]] = {}

    for fname in sorted(os.listdir(kmer_dir)):
        if not fname.lower().endswith(KMER_FILE_EXTS):
            continue

        fpath = os.path.join(kmer_dir, fname)
        if not os.path.isfile(fpath):
            continue

        group = parse_group_from_filename(fname, group_regex)
        group = group.strip()

        group_kmers.setdefault(group, [])

        with open(fpath, "r", encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if (not line) or line.startswith("#"):
                    continue
                if line.lower().startswith(("kmer", "total")):
                    continue

                token = line.split()[0].upper()
                token = clean_protein(token) if is_protein else clean_dna(token)
                if token:
                    group_kmers[group].append(token)

    # Deduplicate while preserving order
    for g in list(group_kmers.keys()):
        group_kmers[g] = list(dict.fromkeys(group_kmers[g]))
        if len(group_kmers[g]) == 0:
            # drop empty groups
            del group_kmers[g]

    if not group_kmers:
        raise FileNotFoundError(f"No k-mer files found in: {kmer_dir}")

    return KmerDB(group_kmers=group_kmers, source_dir=kmer_dir)


def load_unique_kmers(kmer_input: str, is_protein: bool, group_regex: str) -> KmerDB:
    """
    kmer_input can be a directory OR a .zip file containing k-mer output files.
    """
    if os.path.isdir(kmer_input):
        return load_unique_kmers_from_dir(kmer_input, is_protein, group_regex=group_regex)

    if os.path.isfile(kmer_input) and kmer_input.lower().endswith(".zip"):
        tmp = tempfile.mkdtemp(prefix="kmerdb_")
        safe_extract_zip(kmer_input, tmp)
        # find a directory inside that actually contains kmer files; simplest: use tmp itself
        return load_unique_kmers_from_dir(tmp, is_protein, group_regex=group_regex)

    raise FileNotFoundError(f"--kmer-input must be a directory or a .zip file: {kmer_input}")


# ----------------------------
# Matching
# ----------------------------

def align_kmer_to_seq(
    kmer: str,
    seq: str,
    is_protein: bool,
    identity_threshold: float = 0.9,
    min_coverage: float = 0.8,
    gap_open: float = -10,
    gap_extend: float = -0.5,
    nuc_match: float = 2,
    nuc_mismatch: float = -1,
    nuc_gap_open: float = -5,
    nuc_gap_extend: float = -1,
) -> bool:
    if not kmer or not seq:
        return False

    # Fast exact substring path
    if identity_threshold == 1.0 and min_coverage == 1.0:
        return kmer in seq
    if len(kmer) <= 3:
        return kmer in seq

    try:
        aligner = Align.PairwiseAligner()
        if is_protein:
            aligner.substitution_matrix = BLOSUM62
            aligner.open_gap_score = gap_open
            aligner.extend_gap_score = gap_extend
        else:
            aligner.match_score = nuc_match
            aligner.mismatch_score = nuc_mismatch
            aligner.open_gap_score = nuc_gap_open
            aligner.extend_gap_score = nuc_gap_extend

        alns = aligner.align(kmer, seq)
        if not alns:
            return False

        aln = alns[0]
        aligned_query = aln.aligned[0]
        aligned_target = aln.aligned[1]

        aligned_len = sum(e - s for s, e in aligned_query)
        if aligned_len == 0:
            return False

        matches = 0
        for (qs, qe), (ts, te) in zip(aligned_query, aligned_target):
            subseq_q = kmer[qs:qe]
            subseq_t = seq[ts:te]
            matches += sum(1 for a, b in zip(subseq_q, subseq_t) if a == b)

        identity = matches / aligned_len
        coverage = aligned_len / len(kmer)
        return (identity >= identity_threshold) and (coverage >= min_coverage)

    except Exception:
        return False


# ----------------------------
# Prediction core
# ----------------------------

def predict(
    unknown: str,
    kmer_input: str,
    output_dir: str,
    seqtype: str,
    mode: str,
    identity_threshold: float,
    min_coverage: float,
    fdr_alpha: float,
    group_regex: str,
) -> pd.DataFrame:
    is_protein = (seqtype.lower() == "protein")
    mode = mode.lower().strip()
    if mode not in {"fast", "full"}:
        raise ValueError("--mode must be 'fast' or 'full'")

    # Load kmers (dir or zip)
    db = load_unique_kmers(kmer_input, is_protein=is_protein, group_regex=group_regex)
    group_kmers = db.group_kmers

    print(f"Loaded k-mer counts: { {g: len(group_kmers[g]) for g in group_kmers} }")

    # Unknown sequences
    seq_index = iter_unknown_sequences(unknown, is_protein=is_protein)
    if not seq_index:
        raise FileNotFoundError("No sequences found in --unknown (file/dir).")

    # Mode parameters
    if mode == "fast":
        eff_identity = 1.0
        eff_coverage = 1.0
        compute_stats = False
    else:
        eff_identity = float(identity_threshold)
        eff_coverage = float(min_coverage)
        compute_stats = True

    results: List[dict] = []

    total_seqs = len(seq_index)
    for i, (srcfile, header, seq) in enumerate(seq_index, start=1):
        print(f"Processing sequence {i}/{total_seqs} ({os.path.basename(srcfile)})")

        group_found_counts = {g: 0 for g in group_kmers}
        total_found = 0

        for g, kmers in group_kmers.items():
            for kmer in kmers:
                if align_kmer_to_seq(
                    kmer, seq, is_protein=is_protein,
                    identity_threshold=eff_identity,
                    min_coverage=eff_coverage,
                ):
                    group_found_counts[g] += 1
                    total_found += 1

        predicted = max(group_found_counts, key=group_found_counts.get)
        conf_present = (group_found_counts[predicted] / total_found) if total_found else 0.0
        conf_vocab = group_found_counts[predicted] / max(1, len(group_kmers[predicted]))

        row = {
            "Source_file": os.path.basename(srcfile),
            "Sequence": header,
            "Predicted_group": predicted,
            "Matches_total": total_found,
            **{f"Matches_{g}": group_found_counts[g] for g in group_kmers},
            "Confidence_by_present": conf_present,
            "Confidence_by_group_vocab": conf_vocab,
        }

        if compute_stats:
            fisher_p = {}
            # total vocabulary size of "other groups" for contingency table
            other_vocab_total = {g: sum(len(group_kmers[og]) for og in group_kmers if og != g) for g in group_kmers}

            for g in group_kmers:
                a = group_found_counts[g]
                b = len(group_kmers[g]) - a
                c = total_found - a
                d = other_vocab_total[g] - c
                if d < 0:
                    d = 0
                table = [[a, b], [c, d]]
                _, p = fisher_exact(table, alternative="greater")
                fisher_p[g] = p
            row.update({f"FisherP_{g}": fisher_p[g] for g in group_kmers})

        results.append(row)

    df = pd.DataFrame(results)

    # FDR correction (full mode)
    if mode == "full":
        fisher_cols = [c for c in df.columns if c.startswith("FisherP_")]
        if fisher_cols:
            all_pvals = df[fisher_cols].values.flatten()
            _, qvals, _, _ = multipletests(all_pvals, alpha=float(fdr_alpha), method="fdr_bh")
            qval_matrix = qvals.reshape(df[fisher_cols].shape)
            for j, col in enumerate(fisher_cols):
                grp = col.split("_", 1)[1]
                df[f"FDR_{grp}"] = qval_matrix[:, j]

    # Save
    os.makedirs(output_dir, exist_ok=True)
    out_csv = os.path.join(output_dir, "predictions_by_alignment.csv")
    df.to_csv(out_csv, index=False)
    print(f"Saved predictions to {out_csv}")

    # Plot
    save_summary_plot(df, output_dir)

    return df


def save_summary_plot(df: pd.DataFrame, output_dir: str) -> None:
    """
    Matplotlib-only summary figure:
    - Left: predicted group counts
    - Right: confidence distribution (boxplot)
    """
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    # Left: bar counts
    counts = df["Predicted_group"].value_counts()
    axes[0].bar(counts.index.astype(str), counts.values)
    axes[0].set_xlabel("Predicted Group")
    axes[0].set_ylabel("Number of Sequences")
    axes[0].set_title("Predicted Group Counts")
    axes[0].tick_params(axis="x", rotation=45)

    # Right: boxplot confidence_by_present by group
    groups = sorted(df["Predicted_group"].unique().tolist())
    data = [df.loc[df["Predicted_group"] == g, "Confidence_by_present"].values for g in groups]
    axes[1].boxplot(data, labels=groups, showfliers=False)
    axes[1].set_title("Prediction Confidence (by Present)")
    axes[1].set_xlabel("Predicted Group")
    axes[1].set_ylabel("Confidence")
    axes[1].tick_params(axis="x", rotation=45)

    fig.tight_layout()
    fig.savefig(os.path.join(output_dir, "predicted_results_summary.png"), dpi=300)
    plt.close(fig)


# ----------------------------
# CLI
# ----------------------------

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Predict group membership of unknown sequences using unique k-mers.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--unknown", required=True, help="Unknown FASTA file OR directory of FASTA files.")
    p.add_argument("--kmer-input", required=True, help="Directory of unique_k*.tsv/txt OR a ZIP containing them.")
    p.add_argument("--outdir", default="prediction_results", help="Output directory.")
    p.add_argument("--seqtype", choices=["dna", "protein"], default="dna", help="Sequence type.")
    p.add_argument("--mode", choices=["fast", "full"], default="fast", help="fast=substring only; full=alignment+Fisher+FDR.")
    p.add_argument("--identity", type=float, default=0.9, help="Alignment identity threshold (full mode only).")
    p.add_argument("--coverage", type=float, default=0.8, help="Alignment coverage threshold (full mode only).")
    p.add_argument("--fdr", type=float, default=0.05, help="FDR alpha (full mode only).")
    p.add_argument(
        "--group-regex",
        default=DEFAULT_GROUP_REGEX,
        help="Regex to extract group name from k-mer filenames (1st capture group = group).",
    )
    return p


def main() -> None:
    args = build_parser().parse_args()

    # Validate unknown
    if not os.path.exists(args.unknown):
        raise FileNotFoundError(f"--unknown not found: {args.unknown}")

    # Run
    predict(
        unknown=args.unknown,
        kmer_input=args.kmer_input,
        output_dir=args.outdir,
        seqtype=args.seqtype,
        mode=args.mode,
        identity_threshold=args.identity,
        min_coverage=args.coverage,
        fdr_alpha=args.fdr,
        group_regex=args.group_regex,
    )


if __name__ == "__main__":
    main()