"""Map every processed count column to the E-MTAB-5009 metadata."""

from pathlib import Path
import re

import pandas as pd


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DATA_DIR = PROJECT_ROOT / "data" / "raw" / "E-MTAB-5009"
COUNTS_PATH = DATA_DIR / "WP2_100acc.htseq.txt"
METADATA_PATH = DATA_DIR / "E-MTAB-5009.sdrf.txt"
OUTPUT_PATH = PROJECT_ROOT / "data" / "processed" / "sample_sheet.csv"

COUNT_COLUMN_PATTERN = re.compile(
    r"^(?P<ecotype>.+)_(?P<condition>[CS])_B(?P<block>.+?)"
    r"(?:_T(?P<technical_number>\d+))?\.TAIR10_allchr\.htseq$"
)


def normalize_name(value: str) -> str:
    """Make names such as Ag_0 and Ag-0 comparable."""
    return re.sub(r"[^a-z0-9]", "", value.lower())


def technical_sort_key(value: str) -> int:
    match = re.search(r"(\d+)$", value.strip())
    return int(match.group(1)) if match else 1


def parse_count_columns(columns: list[str]) -> pd.DataFrame:
    records = []
    for column in columns:
        match = COUNT_COLUMN_PATTERN.fullmatch(column)
        if match is None:
            raise ValueError(f"Could not parse count column: {column}")

        condition_code = match.group("condition")
        technical_number = match.group("technical_number")
        records.append(
            {
                "count_column": column,
                "normalized_ecotype": normalize_name(match.group("ecotype")),
                "condition": "control" if condition_code == "C" else "drought",
                "block": match.group("block"),
                "column_technical_number": int(technical_number or 1),
            }
        )
    return pd.DataFrame(records)


def prepare_run_metadata() -> pd.DataFrame:
    metadata = pd.read_csv(METADATA_PATH, sep="\t", keep_default_na=False)

    # Paired-end read 1 and read 2 occupy separate SDRF rows but share one run.
    runs = metadata.drop_duplicates(subset="Comment[ENA_RUN]").copy()
    runs["normalized_ecotype"] = runs["Factor Value[ecotype]"].map(normalize_name)
    runs["condition"] = runs["Factor Value[environmental stress]"].map(
        {"control": "control", "Mild drought stress": "drought"}
    )
    runs["block"] = runs["Factor Value[block]"].str.strip()
    runs["metadata_technical_number"] = runs[
        "Characteristics[technical replicate]"
    ].map(technical_sort_key)
    return runs


def build_sample_sheet() -> tuple[pd.DataFrame, pd.DataFrame]:
    count_columns = pd.read_csv(COUNTS_PATH, sep="\t", nrows=0).columns.tolist()
    parsed = parse_count_columns(count_columns)
    runs = prepare_run_metadata()

    keys = ["normalized_ecotype", "condition", "block"]
    mapped_groups = []

    for key, count_group in parsed.groupby(keys, sort=False):
        run_group = runs
        for key_name, key_value in zip(keys, key):
            run_group = run_group.loc[run_group[key_name].eq(key_value)]

        count_group = count_group.sort_values("column_technical_number")
        run_group = run_group.sort_values("metadata_technical_number")

        if len(count_group) != len(run_group):
            raise ValueError(
                f"Mapping mismatch for {key}: "
                f"{len(count_group)} count columns and {len(run_group)} metadata runs"
            )

        count_group = count_group.reset_index(drop=True)
        run_group = run_group.reset_index(drop=True)
        mapped_groups.append(
            pd.DataFrame(
                {
                    "count_column": count_group["count_column"],
                    "sample_group_id": ["|".join(key)] * len(count_group),
                    "ecotype": run_group["Factor Value[ecotype]"],
                    "condition": count_group["condition"],
                    "block": count_group["block"],
                    "technical_replicate": run_group[
                        "Characteristics[technical replicate]"
                    ].str.strip(),
                    "assay_name": run_group["Assay Name"],
                    "ena_sample": run_group["Comment[ENA_SAMPLE]"],
                    "ena_experiment": run_group["Comment[ENA_EXPERIMENT]"],
                    "ena_run": run_group["Comment[ENA_RUN]"],
                }
            )
        )

    sample_sheet = pd.concat(mapped_groups, ignore_index=True)

    processed_keys = set(map(tuple, parsed[keys].drop_duplicates().to_numpy()))
    metadata_groups = runs[keys + ["Factor Value[ecotype]"]].drop_duplicates()
    missing_mask = [
        tuple(row) not in processed_keys
        for row in metadata_groups[keys].to_numpy()
    ]
    missing_processed_groups = metadata_groups.loc[missing_mask].copy()

    return sample_sheet, missing_processed_groups


def validate(sample_sheet: pd.DataFrame) -> None:
    count_columns = pd.read_csv(COUNTS_PATH, sep="\t", nrows=0).columns

    if len(sample_sheet) != len(count_columns):
        raise ValueError("The sample sheet does not contain every count column.")
    if sample_sheet["count_column"].duplicated().any():
        raise ValueError("The sample sheet contains duplicate count columns.")
    if set(sample_sheet["count_column"]) != set(count_columns):
        raise ValueError("Sample-sheet and count-table column names differ.")
    if sample_sheet[
        ["ecotype", "condition", "block", "ena_run"]
    ].eq("").any().any():
        raise ValueError("Required sample-sheet fields contain blank values.")


def main() -> None:
    sample_sheet, missing_groups = build_sample_sheet()
    validate(sample_sheet)

    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    sample_sheet.to_csv(OUTPUT_PATH, index=False)

    group_sizes = sample_sheet.groupby("sample_group_id").size()
    condition_pairs = (
        sample_sheet[["ecotype", "condition"]]
        .drop_duplicates()
        .groupby("ecotype")["condition"]
        .nunique()
    )

    print(f"Wrote: {OUTPUT_PATH}")
    print(f"Mapped count columns: {len(sample_sheet):,}")
    print(f"Sample groups after combining technical replicates: {group_sizes.size:,}")
    print(f"Groups with more than one technical count column: {(group_sizes > 1).sum():,}")
    print(f"Ecotypes in processed counts: {sample_sheet['ecotype'].nunique():,}")
    print(f"Ecotypes having both conditions: {(condition_pairs == 2).sum():,}")
    print("Metadata groups absent from the processed count table:")
    print(
        missing_groups[["Factor Value[ecotype]", "condition", "block"]]
        .to_string(index=False)
    )


if __name__ == "__main__":
    main()