#!/usr/bin/env python3

import argparse
import csv
from collections import Counter
from pathlib import Path


LOCUS_HEADER = [
    "OG",
    "sequence_count",
    "independent_event_count",
    "max_single_sequence_unique_events",
    "carrier_pattern",
    "severe",
]


def read_fasta(fasta_path: Path) -> dict[str, str]:
    sequences: dict[str, list[str]] = {}
    sequence_id: str | None = None

    with fasta_path.open() as fasta:
        for line_number, raw_line in enumerate(fasta, start=1):
            line = raw_line.strip()
            if not line:
                continue
            if line.startswith(">"):
                sequence_id = line[1:].split()[0]
                if not sequence_id:
                    raise ValueError(
                        f"{fasta_path}:{line_number}: empty FASTA identifier"
                    )
                if sequence_id in sequences:
                    raise ValueError(
                        f"{fasta_path}:{line_number}: duplicate FASTA identifier "
                        f"{sequence_id!r}"
                    )
                sequences[sequence_id] = []
            elif sequence_id is None:
                raise ValueError(
                    f"{fasta_path}:{line_number}: sequence before first FASTA header"
                )
            else:
                sequences[sequence_id].append(line)

    if not sequences:
        raise ValueError(f"{fasta_path}: no FASTA sequences found")

    joined = {name: "".join(parts) for name, parts in sequences.items()}
    lengths = {len(sequence) for sequence in joined.values()}
    if len(lengths) != 1:
        raise ValueError(f"{fasta_path}: aligned sequences have different lengths")
    return joined


def event_mode(carrier_count: int, sequence_count: int) -> str:
    if carrier_count == 1:
        return "single_species"
    if carrier_count == sequence_count - 1:
        return "all_but_one"
    if carrier_count == sequence_count:
        return "all_taxa"
    return "branch_shared"


def locus_pattern(
    event_modes: list[str],
    event_count: int,
    max_unique_events: int,
) -> str:
    modes = set(event_modes)
    if max_unique_events >= 2:
        return "single_species_scattered_multiple_events"
    if event_count == 1 and modes == {"single_species"}:
        return "single_species_single_event"
    if modes == {"all_but_one"}:
        return "all_but_one"
    if modes == {"branch_shared"}:
        return "branch_shared"
    if modes == {"all_taxa"}:
        return "all_taxa"
    if not modes:
        return "no_frameshift_detected"
    return "mixed"


def merge_overlapping_triplets(
    triplets: list[tuple[tuple[int, int, int], str]],
) -> list[tuple[tuple[int, ...], list[str]]]:
    components: list[tuple[set[int], set[str]]] = []

    for columns, carrier in triplets:
        column_set = set(columns)
        matching = [
            index
            for index, (component_columns, _) in enumerate(components)
            if component_columns & column_set
        ]

        if not matching:
            components.append((column_set, {carrier}))
            continue

        merged_columns = set(column_set)
        merged_carriers = {carrier}
        for index in reversed(matching):
            component_columns, component_carriers = components.pop(index)
            merged_columns.update(component_columns)
            merged_carriers.update(component_carriers)
        components.append((merged_columns, merged_carriers))

    return [
        (tuple(sorted(columns)), sorted(carriers))
        for columns, carriers in components
    ]


def classify_fasta(fasta_path: Path) -> tuple[list[str], list[list[str]]]:
    sequences = read_fasta(fasta_path)
    sequence_ids = list(sequences)
    sequence_count = len(sequence_ids)
    unique_events_by_sequence: Counter[str] = Counter()
    event_rows: list[list[str]] = []
    modes: list[str] = []
    complete_triplets: list[tuple[tuple[int, int, int], str]] = []
    events: list[tuple[tuple[int, ...], list[str]]] = []

    for sequence_id, sequence in sequences.items():
        marker_columns = [
            column_index + 1
            for column_index, character in enumerate(sequence)
            if character == "!"
        ]
        for index in range(0, len(marker_columns), 3):
            event_columns = tuple(marker_columns[index : index + 3])
            if len(event_columns) == 3:
                complete_triplets.append((event_columns, sequence_id))
            else:
                # A trailing 1-2 markers still form one unique event, but cannot
                # be classified as shared because they are not a complete triplet.
                events.append((event_columns, [sequence_id]))

    events.extend(merge_overlapping_triplets(complete_triplets))
    events.sort(key=lambda event: (event[0], event[1]))

    for event_columns, carriers in events:
        carriers = [
            sequence_id for sequence_id in sequence_ids if sequence_id in carriers
        ]
        mode = event_mode(len(carriers), sequence_count)
        modes.append(mode)
        if len(carriers) == 1:
            unique_events_by_sequence[carriers[0]] += 1
        event_rows.append(
            [
                fasta_path.stem,
                ",".join(str(column) for column in event_columns),
                str(len(carriers)),
                ",".join(carriers),
                mode,
            ]
        )

    event_count = len(event_rows)
    max_unique_events = max(unique_events_by_sequence.values(), default=0)
    severe = event_count >= 3 or max_unique_events >= 2
    locus_row = [
        fasta_path.stem,
        str(sequence_count),
        str(event_count),
        str(max_unique_events),
        locus_pattern(modes, event_count, max_unique_events),
        "yes" if severe else "no",
    ]
    return locus_row, event_rows


def fasta_paths_from_list(list_path: Path) -> list[Path]:
    paths: list[Path] = []
    seen: set[Path] = set()
    with list_path.open() as fasta_list:
        for line_number, raw_line in enumerate(fasta_list, start=1):
            value = raw_line.strip()
            if not value:
                continue
            path = Path(value)
            if not path.is_absolute():
                path = list_path.parent / path
            path = path.resolve()
            if not path.is_file():
                raise FileNotFoundError(
                    f"{list_path}:{line_number}: FASTA file not found: {path}"
                )
            if path not in seen:
                paths.append(path)
                seen.add(path)
    return paths


def write_results(
    output_dir: Path,
    locus_rows: list[list[str]],
    event_rows: list[list[str]],
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)

    with (output_dir / "frameshift_events.tsv").open("w", newline="") as output:
        writer = csv.writer(output, delimiter="\t", lineterminator="\n")
        writer.writerow(["OG", "alignment_column", "carrier_count", "carriers", "mode"])
        writer.writerows(event_rows)

    with (output_dir / "frameshift_locus_classification.tsv").open(
        "w", newline=""
    ) as output:
        writer = csv.writer(output, delimiter="\t", lineterminator="\n")
        writer.writerow(LOCUS_HEADER)
        writer.writerows(locus_rows)

    severe_rows = [row for row in locus_rows if row[-1] == "yes"]
    with (output_dir / "frameshift_severe_subset.tsv").open(
        "w", newline=""
    ) as output:
        writer = csv.writer(output, delimiter="\t", lineterminator="\n")
        writer.writerow(LOCUS_HEADER)
        writer.writerows(severe_rows)

    pattern_counts = Counter(row[4] for row in locus_rows)
    with (output_dir / "frameshift_classification_summary.tsv").open(
        "w", newline=""
    ) as output:
        writer = csv.writer(output, delimiter="\t", lineterminator="\n")
        writer.writerow(["metric", "count"])
        writer.writerow(["total_FS_loci", len(locus_rows)])
        writer.writerow(["severe_loci", len(severe_rows)])
        for pattern, count in sorted(pattern_counts.items()):
            writer.writerow([pattern, count])


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Classify frameshift loci by event count and carrier pattern."
    )
    parser.add_argument(
        "--fasta-list",
        required=True,
        type=Path,
        help="File containing one aligned FASTA path per line.",
    )
    parser.add_argument(
        "--output-dir",
        required=True,
        type=Path,
        help="Directory for classification tables.",
    )
    args = parser.parse_args()

    fasta_paths = fasta_paths_from_list(args.fasta_list.resolve())
    locus_rows: list[list[str]] = []
    event_rows: list[list[str]] = []
    for fasta_path in fasta_paths:
        locus_row, rows = classify_fasta(fasta_path)
        locus_rows.append(locus_row)
        event_rows.extend(rows)
    write_results(args.output_dir, locus_rows, event_rows)


if __name__ == "__main__":
    main()