"""Prepare a safe classroom feature table from the three UCI CSV files."""

from pathlib import Path
import re

import numpy as np
import pandas as pd


HERE = Path(__file__).resolve().parent
RAW_DIR = HERE / "raw"
OUTPUT_PATH = HERE / "malware_features.csv"
MANIFEST_PATH = HERE / "feature_manifest.csv"

FILES = {
    "benign_lab": ("staDynBenignLab.csv", 0),
    "vxheaven": ("staDynVxHeaven2698Lab.csv", 1),
    "virustotal": ("staDynVt2955Lab.csv", 1),
}

FORBIDDEN_NAMES = {
    "label", "class", "target", "malware", "is_malware", "type",
    "file", "filename", "file_name", "name", "hash", "md5", "sha1",
    "sha256", "source", "category",
}


def normalize_name(value: object) -> str:
    """Convert a source column name to a stable snake_case identifier."""
    text = re.sub(r"[^0-9A-Za-z]+", "_", str(value).strip()).strip("_").lower()
    return text or "unnamed_feature"


def normalize_columns(columns: list[object]) -> tuple[list[str], dict[str, list[str]]]:
    """Normalize names and disambiguate collisions by their source-column order.

    The UCI files intentionally contain case- or punctuation-distinct fields such as
    ``WORD``/``word`` and ``Import``/``__IMPORT_``. A simple lowercase snake-case
    conversion collapses each pair. The three source files use the same relative
    order for these fields, so an occurrence suffix produces stable shared names:
    ``word``, ``word__2``, and so on.
    """
    counts: dict[str, int] = {}
    normalized: list[str] = []
    originals_by_base: dict[str, list[str]] = {}

    for raw_column in columns:
        base = normalize_name(raw_column)
        counts[base] = counts.get(base, 0) + 1
        occurrence = counts[base]
        unique_name = base if occurrence == 1 else f"{base}__{occurrence}"
        normalized.append(unique_name)
        originals_by_base.setdefault(base, []).append(str(raw_column))

    collisions = {
        base: originals
        for base, originals in originals_by_base.items()
        if len(originals) > 1
    }
    if len(normalized) != len(set(normalized)):
        raise ValueError("Column names are still duplicated after disambiguation.")
    return normalized, collisions


def find_required_file(filename: str) -> Path:
    """Find one required CSV below data/raw, allowing an extracted subfolder."""
    matches = list(RAW_DIR.rglob(filename)) if RAW_DIR.exists() else []
    if len(matches) != 1:
        raise FileNotFoundError(
            f"Expected exactly one {filename!r} below {RAW_DIR}. "
            "Follow data/README.md and check the extracted folder structure."
        )
    return matches[0]


def load_source(source: str, filename: str, label: int) -> pd.DataFrame:
    """Load one source file and attach provenance columns."""
    frame = pd.read_csv(find_required_file(filename), low_memory=False)
    normalized, collisions = normalize_columns(frame.columns.tolist())
    frame.columns = normalized
    # Consolidate the wide table before adding provenance columns. This avoids
    # pandas fragmentation warnings on the 1,000+ column UCI source files.
    frame = frame.copy()
    frame.attrs["normalized_column_collisions"] = collisions
    frame["source"] = source
    frame["label"] = label
    return frame


def validate_collision_alignment(frames: dict[str, pd.DataFrame]) -> None:
    """Ensure occurrence suffixes refer to the same raw fields in every source."""
    collision_names = set().union(*(
        set(frame.attrs.get("normalized_column_collisions", {}))
        for frame in frames.values()
    ))
    for base in sorted(collision_names):
        variants = {
            source: tuple(frame.attrs.get("normalized_column_collisions", {}).get(base, []))
            for source, frame in frames.items()
        }
        if len(set(variants.values())) != 1:
            raise ValueError(
                f"Normalized collision {base!r} is not aligned across source files: {variants}"
            )

    if collision_names:
        print("Disambiguated normalized column collisions:")
        reference = next(iter(frames.values())).attrs["normalized_column_collisions"]
        for base in sorted(collision_names):
            generated = [base] + [f"{base}__{index}" for index in range(2, len(reference[base]) + 1)]
            print(f"  {reference[base]} -> {generated}")


def main() -> None:
    frames = {
        source: load_source(source, filename, label)
        for source, (filename, label) in FILES.items()
    }
    validate_collision_alignment(frames)

    common = set.intersection(*(set(frame.columns) for frame in frames.values()))
    candidate_features = sorted(common - FORBIDDEN_NAMES - {"source", "label"})
    if not candidate_features:
        raise ValueError("No shared candidate feature columns were found.")

    combined_parts = []
    for source, frame in frames.items():
        part = frame[candidate_features].copy()
        part.insert(0, "source", source)
        part.insert(1, "label", int(FILES[source][1]))
        combined_parts.append(part)
    combined = pd.concat(combined_parts, ignore_index=True)

    numeric = combined[candidate_features].apply(pd.to_numeric, errors="coerce")
    numeric = numeric.replace([np.inf, -np.inf], np.nan)
    usable_ratio = numeric.notna().mean()
    retained = usable_ratio[usable_ratio >= 0.95].index.tolist()
    numeric = numeric[retained]

    non_constant = numeric.nunique(dropna=True)
    retained = non_constant[non_constant > 1].index.tolist()
    numeric = numeric[retained]
    if numeric.empty:
        raise ValueError("No usable numeric features remained after validation.")

    prepared = pd.concat([combined[["source", "label"]], numeric], axis=1)
    prepared = prepared.sample(frac=1.0, random_state=42).reset_index(drop=True)
    prepared.to_csv(OUTPUT_PATH, index=False)

    manifest = pd.DataFrame({
        "feature": retained,
        "non_missing_ratio": numeric[retained].notna().mean().values,
        "unique_values": numeric[retained].nunique(dropna=True).values,
        "feature_family": "unclassified",
    })
    manifest.to_csv(MANIFEST_PATH, index=False)

    print(f"Created {OUTPUT_PATH}")
    print(f"Shape: {prepared.shape}")
    print("Class counts:")
    print(prepared["label"].value_counts().sort_index())
    print("Source counts:")
    print(prepared["source"].value_counts())
    print(f"Created {MANIFEST_PATH}")


if __name__ == "__main__":
    main()