# ABOUTME: Core implementation of MSA summary statistics calculator.
# ABOUTME: Computes 27 metrics matching the original C++ msastats behavior.
"""MSA Statistics Calculator.

A Python implementation of the MsaStatsCalculator for computing MSA summary statistics.
"""
import enum
from pathlib import Path
from typing import List, Dict, Tuple, Union


class MsaStatsError(Exception):
    """Base exception for MSA statistics errors."""
    pass

class StatType(enum.Enum):
    """Defines the types of summary statistics that can be calculated."""
    AVG_GAP_SIZE = enum.auto()
    MSA_LEN = enum.auto()
    LONGEST_UNALIGNED_SEQ = enum.auto()
    SHORTEST_UNALIGNED_SEQ = enum.auto()
    TOT_NUM_GAPS = enum.auto()
    NUM_GAPS_LEN_ONE = enum.auto()
    NUM_GAPS_LEN_TWO = enum.auto()
    NUM_GAPS_LEN_THREE = enum.auto()
    NUM_GAPS_LEN_AT_LEAST_FOUR = enum.auto()
    AVG_UNIQUE_GAP_SIZE = enum.auto()
    TOT_NUM_UNIQUE_GAPS = enum.auto()
    NUM_GAPS_LEN_ONE_IN_ONE_SEQ = enum.auto()
    NUM_GAPS_LEN_ONE_IN_TWO_SEQS = enum.auto()
    NUM_GAPS_LEN_ONE_IN_ALL_EXCEPT_ONE = enum.auto()
    NUM_GAPS_LEN_TWO_IN_ONE_SEQ = enum.auto()
    NUM_GAPS_LEN_TWO_IN_TWO_SEQS = enum.auto()
    NUM_GAPS_LEN_TWO_IN_ALL_EXCEPT_ONE = enum.auto()
    NUM_GAPS_LEN_THREE_IN_ONE_SEQ = enum.auto()
    NUM_GAPS_LEN_THREE_IN_TWO_SEQS = enum.auto()
    NUM_GAPS_LEN_THREE_IN_ALL_EXCEPT_ONE = enum.auto()
    NUM_GAPS_LEN_AT_LEAST_FOUR_IN_ONE_SEQ = enum.auto()
    NUM_GAPS_LEN_AT_LEAST_FOUR_IN_TWO_SEQS = enum.auto()
    NUM_GAPS_LEN_AT_LEAST_FOUR_IN_ALL_EXCEPT_ONE = enum.auto()
    MSA_POSITION_WITH_0_GAPS = enum.auto()
    MSA_POSITION_WITH_1_GAPS = enum.auto()
    MSA_POSITION_WITH_2_GAPS = enum.auto()
    MSA_POSITION_WITH_N_MINUS_1_GAPS = enum.auto()


class MsaStatsCalculator:
    """
    Calculates a variety of summary statistics for a given Multiple Sequence Alignment (MSA).
    The algorithms and data structures are designed to mirror the original C++ implementation.
    """

    def __init__(self, msa_sequences: List[str]):
        """Initializes the calculator with an MSA provided as a list of strings.

        Args:
            msa_sequences: A list of strings, where each string is an aligned sequence.

        Raises:
            ValueError: If input is empty, contains non-strings, has empty sequences,
                or sequences have mismatched lengths.
        """
        if not msa_sequences:
            raise ValueError("Input must be a non-empty list of sequences.")

        for i, seq in enumerate(msa_sequences):
            if not isinstance(seq, str):
                raise ValueError(
                    f"Sequence at index {i} is not a string (got {type(seq).__name__})."
                )

        self._original_aligned_seqs: Tuple[str, ...] = tuple(msa_sequences)
        self._number_of_sequences: int = len(self._original_aligned_seqs)

        if not self._original_aligned_seqs[0]:
            raise ValueError("Sequences cannot be empty.")

        self._msa_length = len(self._original_aligned_seqs[0])
        for i, seq in enumerate(self._original_aligned_seqs[1:], start=1):
            if len(seq) != self._msa_length:
                raise ValueError(
                    f"Sequence length mismatch: sequence 0 has length {self._msa_length}, "
                    f"but sequence {i} has length {len(seq)}."
                )

        self._initialize_all_variables()

    @classmethod
    def from_fasta(cls, fasta_path: Union[str, Path]) -> 'MsaStatsCalculator':
        """Creates an MsaStatsCalculator instance from a FASTA file.

        Args:
            fasta_path: The path to the FASTA file.

        Returns:
            An instance of MsaStatsCalculator.

        Raises:
            MsaStatsError: If the file cannot be read or contains no sequences.
            ValueError: If the parsed sequences are invalid (empty, mismatched lengths).
        """
        fasta_path = Path(fasta_path)
        try:
            with open(fasta_path, 'r') as f:
                sequences = []
                current_seq = ""
                for line in f:
                    line = line.strip()
                    if line.startswith('>'):
                        if current_seq:
                            sequences.append(current_seq)
                        current_seq = ""
                    else:
                        current_seq += line
                if current_seq:
                    sequences.append(current_seq)
        except FileNotFoundError as e:
            raise MsaStatsError(f"FASTA file not found: {fasta_path}") from e
        except PermissionError as e:
            raise MsaStatsError(f"Permission denied reading FASTA file: {fasta_path}") from e
        except IsADirectoryError as e:
            raise MsaStatsError(f"Path is a directory, not a file: {fasta_path}") from e
        except OSError as e:
            raise MsaStatsError(f"Failed to read FASTA file '{fasta_path}': {e}") from e

        if not sequences:
            raise MsaStatsError(f"No sequences found in FASTA file: {fasta_path}")

        return cls(sequences)

    def _initialize_all_variables(self) -> None:
        """Resets all internal statistics and working data structures."""
        self._aligned_seqs: List[str] = list(self._original_aligned_seqs)

        # Core data structures
        self._unique_indel_map: Dict[Tuple[int, int], List[int]] = {}
        self._indel_counter: List[int] = []

        # Final statistics attributes
        self._ave_indel_length: float = 0.0
        self._total_number_of_indels: int = 0
        self._longest_seq_length: int = 0
        self._shortest_seq_length: int = 0

        self._number_of_indels_of_length_one: int = 0
        self._number_of_indels_of_length_two: int = 0
        self._number_of_indels_of_length_three: int = 0
        self._number_of_indels_of_length_at_least_four: int = 0

        self._number_of_indels_of_length_one_in_one_position: int = 0
        self._number_of_indels_of_length_one_in_two_positions: int = 0
        self._number_of_indels_of_length_one_in_n_minus_1_positions: int = 0

        self._number_of_indels_of_length_two_in_one_position: int = 0
        self._number_of_indels_of_length_two_in_two_positions: int = 0
        self._number_of_indels_of_length_two_in_n_minus_1_positions: int = 0

        self._number_of_indels_of_length_three_in_one_position: int = 0
        self._number_of_indels_of_length_three_in_two_positions: int = 0
        self._number_of_indels_of_length_three_in_n_minus_1_positions: int = 0

        self._number_of_indels_of_length_at_least_four_in_one_position: int = 0
        self._number_of_indels_of_length_at_least_four_in_two_positions: int = 0
        self._number_of_indels_of_length_at_least_four_in_n_minus_1_positions: int = 0

        self._number_of_msa_position_with_0_gaps: int = 0
        self._number_of_msa_position_with_1_gaps: int = 0
        self._number_of_msa_position_with_2_gaps: int = 0
        self._number_of_msa_position_with_n_minus_1_gaps: int = 0

        # Unique indels summary statistics
        self._ave_unique_indel_length: float = 0.0
        self._total_number_of_unique_indels: int = 0

    def _trim_msa_from_all_indel_position_and_get_summary_statistics_from_indel_counter(self) -> None:
        """
        Counts gaps per column, calculates column-based stats, and trims all-gap columns.
        This method modifies self._aligned_seqs.
        """
        if not self._original_aligned_seqs:
            return

        self._indel_counter = [0] * self._msa_length
        for seq in self._original_aligned_seqs:
            for i, char in enumerate(seq):
                if char == '-':
                    self._indel_counter[i] += 1

        for count in self._indel_counter:
            if count == 0:
                self._number_of_msa_position_with_0_gaps += 1
            elif count == 1:
                self._number_of_msa_position_with_1_gaps += 1
            elif count == 2:
                self._number_of_msa_position_with_2_gaps += 1
            elif count == self._number_of_sequences - 1:
                self._number_of_msa_position_with_n_minus_1_gaps += 1
        
        # Identify columns to keep
        cols_to_keep = [i for i, count in enumerate(self._indel_counter) if count < self._number_of_sequences]

        # Create new sequences with only the columns to keep
        self._aligned_seqs = [''.join(seq[i] for i in cols_to_keep) for seq in self._original_aligned_seqs]


    def _fill_unique_gaps_map(self) -> None:
        """
        Scans the aligned sequences to identify and count unique indel events,
        populating the _unique_indel_map.
        """
        self._unique_indel_map.clear()
        if not self._aligned_seqs:
            return
            
        msa_length = len(self._aligned_seqs[0])

        for seq in self._aligned_seqs:
            in_indel = False
            start_index = -1

            for i, char in enumerate(seq):
                if char == '-' and not in_indel:
                    in_indel = True
                    start_index = i
                elif char != '-' and in_indel:
                    # End of the indel, record it
                    end_index = i - 1
                    key = (start_index, end_index)
                    length = end_index - start_index + 1
                    
                    if key not in self._unique_indel_map:
                        self._unique_indel_map[key] = [length, 0]
                    self._unique_indel_map[key][1] += 1
                    
                    in_indel = False
                    start_index = -1

            # Handle indel that goes to the end of the sequence
            if in_indel:
                end_index = msa_length - 1
                key = (start_index, end_index)
                length = end_index - start_index + 1
                if key not in self._unique_indel_map:
                    self._unique_indel_map[key] = [length, 0]
                self._unique_indel_map[key][1] += 1

    def _set_values_of_indel_summ_stats(self) -> None:
        """
        Calculates indel-related summary statistics by processing the _unique_indel_map.
        """
        self._fill_unique_gaps_map()

        total_gap_chars = 0
        total_unique_gap_chars = 0

        for (length, count) in self._unique_indel_map.values():
            self._total_number_of_indels += count
            self._total_number_of_unique_indels += 1
            total_gap_chars += length * count
            total_unique_gap_chars += length

            if length == 1:
                self._number_of_indels_of_length_one += count
                if count == 1:
                    self._number_of_indels_of_length_one_in_one_position += 1
                if count == 2:
                    self._number_of_indels_of_length_one_in_two_positions += 1
                if count == self._number_of_sequences - 1:
                    self._number_of_indels_of_length_one_in_n_minus_1_positions += 1
            elif length == 2:
                self._number_of_indels_of_length_two += count
                if count == 1:
                    self._number_of_indels_of_length_two_in_one_position += 1
                if count == 2:
                    self._number_of_indels_of_length_two_in_two_positions += 1
                if count == self._number_of_sequences - 1:
                    self._number_of_indels_of_length_two_in_n_minus_1_positions += 1
            elif length == 3:
                self._number_of_indels_of_length_three += count
                if count == 1:
                    self._number_of_indels_of_length_three_in_one_position += 1
                if count == 2:
                    self._number_of_indels_of_length_three_in_two_positions += 1
                if count == self._number_of_sequences - 1:
                    self._number_of_indels_of_length_three_in_n_minus_1_positions += 1
            else: # length >= 4
                self._number_of_indels_of_length_at_least_four += count
                if count == 1:
                    self._number_of_indels_of_length_at_least_four_in_one_position += 1
                if count == 2:
                    self._number_of_indels_of_length_at_least_four_in_two_positions += 1
                if count == self._number_of_sequences - 1:
                    self._number_of_indels_of_length_at_least_four_in_n_minus_1_positions += 1

        if self._total_number_of_indels > 0:
            self._ave_indel_length = total_gap_chars / self._total_number_of_indels
        if self._total_number_of_unique_indels > 0:
            self._ave_unique_indel_length = total_unique_gap_chars / self._total_number_of_unique_indels

    def _set_longest_and_shortest_sequence_lengths(self) -> None:
        """Calculates the longest and shortest ungapped sequence lengths."""
        if not self._original_aligned_seqs:
            self._longest_seq_length = 0
            self._shortest_seq_length = 0
            return

        seq_lengths = [len(s.replace('-', '')) for s in self._original_aligned_seqs]
        self._longest_seq_length = max(seq_lengths)
        self._shortest_seq_length = min(seq_lengths)


    def recompute_stats(self) -> None:
        """
        The main public method to run the full statistical analysis.
        This orchestrates the calls to the internal calculation methods in the correct order.
        """
        self._initialize_all_variables()
        self._trim_msa_from_all_indel_position_and_get_summary_statistics_from_indel_counter()
        self._set_values_of_indel_summ_stats()
        self._set_longest_and_shortest_sequence_lengths()

    def get_stat_by_type(self, stat_type: StatType) -> float:
        """
        Returns the value of a specific statistic.

        Args:
            stat_type: The enum member representing the statistic to retrieve.

        Returns:
            The calculated value of the statistic as a float.
        """
        stat_map = {
            StatType.AVG_GAP_SIZE: self.average_indel_size,
            StatType.MSA_LEN: self.msa_length,
            StatType.LONGEST_UNALIGNED_SEQ: self.msa_longest_seq_length,
            StatType.SHORTEST_UNALIGNED_SEQ: self.msa_shortest_seq_length,
            StatType.TOT_NUM_GAPS: self.total_number_of_indels,
            StatType.NUM_GAPS_LEN_ONE: self.number_of_indels_of_length_one,
            StatType.NUM_GAPS_LEN_TWO: self.number_of_indels_of_length_two,
            StatType.NUM_GAPS_LEN_THREE: self.number_of_indels_of_length_three,
            StatType.NUM_GAPS_LEN_AT_LEAST_FOUR: self.number_of_indels_of_length_at_least_four,
            StatType.AVG_UNIQUE_GAP_SIZE: self.average_unique_indel_size,
            StatType.TOT_NUM_UNIQUE_GAPS: self.total_number_of_unique_indels,
            StatType.NUM_GAPS_LEN_ONE_IN_ONE_SEQ: self.number_of_indels_of_length_one_in_one_position,
            StatType.NUM_GAPS_LEN_ONE_IN_TWO_SEQS: self.number_of_indels_of_length_one_in_two_positions,
            StatType.NUM_GAPS_LEN_ONE_IN_ALL_EXCEPT_ONE: self.number_of_indels_of_length_one_in_n_minus_1_positions,
            StatType.NUM_GAPS_LEN_TWO_IN_ONE_SEQ: self.number_of_indels_of_length_two_in_one_position,
            StatType.NUM_GAPS_LEN_TWO_IN_TWO_SEQS: self.number_of_indels_of_length_two_in_two_positions,
            StatType.NUM_GAPS_LEN_TWO_IN_ALL_EXCEPT_ONE: self.number_of_indels_of_length_two_in_n_minus_1_positions,
            StatType.NUM_GAPS_LEN_THREE_IN_ONE_SEQ: self.number_of_indels_of_length_three_in_one_position,
            StatType.NUM_GAPS_LEN_THREE_IN_TWO_SEQS: self.number_of_indels_of_length_three_in_two_positions,
            StatType.NUM_GAPS_LEN_THREE_IN_ALL_EXCEPT_ONE: self.number_of_indels_of_length_three_in_n_minus_1_positions,
            StatType.NUM_GAPS_LEN_AT_LEAST_FOUR_IN_ONE_SEQ: self.number_of_indels_of_length_at_least_four_in_one_position,
            StatType.NUM_GAPS_LEN_AT_LEAST_FOUR_IN_TWO_SEQS: self.number_of_indels_of_length_at_least_four_in_two_positions,
            StatType.NUM_GAPS_LEN_AT_LEAST_FOUR_IN_ALL_EXCEPT_ONE: self.number_of_indels_of_length_at_least_four_in_n_minus_1_positions,
            StatType.MSA_POSITION_WITH_0_GAPS: self.number_of_msa_position_with_0_gaps,
            StatType.MSA_POSITION_WITH_1_GAPS: self.number_of_msa_position_with_1_gaps,
            StatType.MSA_POSITION_WITH_2_GAPS: self.number_of_msa_position_with_2_gaps,
            StatType.MSA_POSITION_WITH_N_MINUS_1_GAPS: self.number_of_msa_position_with_n_minus_1_gaps,
        }
        return float(stat_map.get(stat_type, -1.0))

    def get_stat_vec(self) -> List[float]:
        """
        Returns a list of all summary statistics in a predefined order, matching
        the C++ implementation's `getStatVec` method.
        """
        return [self.get_stat_by_type(stat) for stat in StatType]

    def __str__(self) -> str:
        """Returns a string representation of the original aligned MSA."""
        return "\n".join(self._original_aligned_seqs)

    # Public properties for accessing statistics
    @property
    def msa_length(self) -> int:
        return self._msa_length

    @property
    def number_of_sequences(self) -> int:
        return self._number_of_sequences

    @property
    def total_number_of_indels(self) -> int:
        return self._total_number_of_indels

    @property
    def total_number_of_unique_indels(self) -> int:
        return self._total_number_of_unique_indels

    @property
    def number_of_indels_of_length_one(self) -> int:
        return self._number_of_indels_of_length_one

    @property
    def number_of_indels_of_length_two(self) -> int:
        return self._number_of_indels_of_length_two

    @property
    def number_of_indels_of_length_three(self) -> int:
        return self._number_of_indels_of_length_three

    @property
    def number_of_indels_of_length_at_least_four(self) -> int:
        return self._number_of_indels_of_length_at_least_four

    @property
    def average_indel_size(self) -> float:
        return self._ave_indel_length

    @property
    def average_unique_indel_size(self) -> float:
        return self._ave_unique_indel_length

    @property
    def msa_longest_seq_length(self) -> int:
        return self._longest_seq_length

    @property
    def msa_shortest_seq_length(self) -> int:
        return self._shortest_seq_length

    @property
    def number_of_indels_of_length_one_in_one_position(self) -> int:
        return self._number_of_indels_of_length_one_in_one_position

    @property
    def number_of_indels_of_length_one_in_two_positions(self) -> int:
        return self._number_of_indels_of_length_one_in_two_positions

    @property
    def number_of_indels_of_length_one_in_n_minus_1_positions(self) -> int:
        return self._number_of_indels_of_length_one_in_n_minus_1_positions

    @property
    def number_of_indels_of_length_two_in_one_position(self) -> int:
        return self._number_of_indels_of_length_two_in_one_position

    @property
    def number_of_indels_of_length_two_in_two_positions(self) -> int:
        return self._number_of_indels_of_length_two_in_two_positions

    @property
    def number_of_indels_of_length_two_in_n_minus_1_positions(self) -> int:
        return self._number_of_indels_of_length_two_in_n_minus_1_positions

    @property
    def number_of_indels_of_length_three_in_one_position(self) -> int:
        return self._number_of_indels_of_length_three_in_one_position

    @property
    def number_of_indels_of_length_three_in_two_positions(self) -> int:
        return self._number_of_indels_of_length_three_in_two_positions

    @property
    def number_of_indels_of_length_three_in_n_minus_1_positions(self) -> int:
        return self._number_of_indels_of_length_three_in_n_minus_1_positions

    @property
    def number_of_indels_of_length_at_least_four_in_one_position(self) -> int:
        return self._number_of_indels_of_length_at_least_four_in_one_position

    @property
    def number_of_indels_of_length_at_least_four_in_two_positions(self) -> int:
        return self._number_of_indels_of_length_at_least_four_in_two_positions

    @property
    def number_of_indels_of_length_at_least_four_in_n_minus_1_positions(self) -> int:
        return self._number_of_indels_of_length_at_least_four_in_n_minus_1_positions

    @property
    def number_of_msa_position_with_0_gaps(self) -> int:
        return self._number_of_msa_position_with_0_gaps

    @property
    def number_of_msa_position_with_1_gaps(self) -> int:
        return self._number_of_msa_position_with_1_gaps

    @property
    def number_of_msa_position_with_2_gaps(self) -> int:
        return self._number_of_msa_position_with_2_gaps

    @property
    def number_of_msa_position_with_n_minus_1_gaps(self) -> int:
        return self._number_of_msa_position_with_n_minus_1_gaps
