"""Combine technical counts, filter genes, and normalize E-MTAB-5009."""

from math import ceil
from pathlib import Path

import numpy as np
import pandas as pd


PROJECT_ROOT = Path(__file__).resolve().parents[1]
RAW_COUNTS_PATH = (
    PROJECT_ROOT / "data" / "raw" / "E-MTAB-5009" / "WP2_100acc.htseq.txt"
)
SAMPLE_SHEET_PATH = PROJECT_ROOT / "data" / "processed" / "sample_sheet.csv"
OUTPUT_DIR = PROJECT_ROOT / "data" / "processed"

MIN_CPM = 1.0
MIN_SAMPLE_FRACTION = 0.10


def load_inputs() -> tuple[pd.DataFrame, pd.DataFrame]:
    counts = pd.read_csv(RAW_COUNTS_PATH, sep="\t", index_col=0)
    counts.index.name = "gene_id"
    sample_sheet = pd.read_csv(SAMPLE_SHEET_PATH, keep_default_na=False)

    if counts.columns.duplicated().any():
        raise ValueError("Raw count columns are not unique.")
    if sample_sheet["count_column"].duplicated().any():
        raise ValueError("Sample-sheet count columns are not unique.")
    if set(counts.columns) != set(sample_sheet["count_column"]):
        raise ValueError("Raw counts and sample sheet contain different columns.")
    if (counts < 0).any().any():
        raise ValueError("Raw counts contain negative values.")

    # Put sample-sheet rows in exactly the same order as the count columns.
    sample_sheet = sample_sheet.set_index("count_column").loc[counts.columns]
    sample_sheet.index.name = "count_column"
    sample_sheet = sample_sheet.reset_index()
    return counts, sample_sheet


def combine_technical_counts(
    counts: pd.DataFrame, sample_sheet: pd.DataFrame
) -> tuple[pd.DataFrame, pd.DataFrame]:
    column_to_group = sample_sheet.set_index("count_column")["sample_group_id"]

    # Raw sequencing counts from technical measurements are additive.
    aggregated = (
        counts.T.join(column_to_group)
        .groupby("sample_group_id", sort=False)
        .sum()
        .T
    )
    aggregated.index.name = "gene_id"

    group_sheet = (
        sample_sheet.groupby("sample_group_id", sort=False)
        .agg(
            ecotype=("ecotype", "first"),
            condition=("condition", "first"),
            block=("block", "first"),
            technical_count_columns=("count_column", "count"),
            source_count_columns=("count_column", ";".join),
        )
        .reset_index()
    )

    if list(aggregated.columns) != group_sheet["sample_group_id"].tolist():
        raise ValueError("Aggregated counts and group sheet have different orders.")

    return aggregated, group_sheet


def normalize_and_filter(
    aggregated: pd.DataFrame,
) -> tuple[pd.DataFrame, pd.DataFrame, int]:
    library_sizes = aggregated.sum(axis=0)
    if (library_sizes == 0).any():
        raise ValueError("At least one sample group has a library size of zero.")

    cpm = aggregated.div(library_sizes, axis="columns") * 1_000_000
    minimum_samples = ceil(MIN_SAMPLE_FRACTION * aggregated.shape[1])
    keep_gene = (cpm >= MIN_CPM).sum(axis=1) >= minimum_samples

    filtered_counts = aggregated.loc[keep_gene]
    log2_cpm = np.log2(cpm.loc[keep_gene] + 1).astype("float32")
    return filtered_counts, log2_cpm, minimum_samples


def main() -> None:
    counts, sample_sheet = load_inputs()
    aggregated, group_sheet = combine_technical_counts(counts, sample_sheet)
    filtered_counts, log2_cpm, minimum_samples = normalize_and_filter(aggregated)

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    aggregated.to_csv(
        OUTPUT_DIR / "counts_aggregated.tsv.gz", sep="\t", compression="gzip"
    )
    filtered_counts.to_csv(
        OUTPUT_DIR / "counts_filtered.tsv.gz", sep="\t", compression="gzip"
    )
    log2_cpm.to_csv(
        OUTPUT_DIR / "expression_log2_cpm.tsv.gz", sep="\t", compression="gzip"
    )
    group_sheet.to_csv(OUTPUT_DIR / "sample_groups.csv", index=False)

    original_libraries = counts.sum(axis=0)
    combined_libraries = aggregated.sum(axis=0)
    print(f"Original count columns: {counts.shape[1]:,}")
    print(f"Sample groups after technical summation: {aggregated.shape[1]:,}")
    print(f"Genes before filtering: {aggregated.shape[0]:,}")
    print(
        f"Filter: CPM >= {MIN_CPM:g} in at least "
        f"{minimum_samples} of {aggregated.shape[1]} sample groups"
    )
    print(f"Genes retained: {filtered_counts.shape[0]:,}")
    print(f"Genes removed: {aggregated.shape[0] - filtered_counts.shape[0]:,}")
    print(
        "Original library sizes (min / median / max): "
        f"{original_libraries.min():,.0f} / "
        f"{original_libraries.median():,.0f} / "
        f"{original_libraries.max():,.0f}"
    )
    print(
        "Combined library sizes (min / median / max): "
        f"{combined_libraries.min():,.0f} / "
        f"{combined_libraries.median():,.0f} / "
        f"{combined_libraries.max():,.0f}"
    )
    print("Wrote counts_aggregated.tsv.gz")
    print("Wrote counts_filtered.tsv.gz")
    print("Wrote expression_log2_cpm.tsv.gz")
    print("Wrote sample_groups.csv")


if __name__ == "__main__":
    main()