"""Efficient mutation sequence generation module."""

import itertools as it
from copy import deepcopy
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

from kdpred.utils import STANDARD_AMINO_ACIDS, ensure_directory, get_chain_index, parse_mutation


class MutationGenerator:
    """Efficient generator for creating mutated protein sequences."""

    def __init__(self, protein_sequence: List[str], protein_name: str = ""):
        """Initialize with wildtype protein sequences.

        Parameters
        ----------
        protein_sequence : List[str]
            List of protein sequences, one per chain.
        protein_name : str, optional
            Name of the protein, by default "".

        """
        self.protein_sequence = [seq.upper() for seq in protein_sequence]
        self.num_chains = len(protein_sequence)
        self.protein_name = protein_name

    def _apply_mutation_to_sequence(self, sequence: str, position: int, mutant: str) -> str:
        """Apply a single mutation to a sequence string.

        Parameters
        ----------
        sequence : str
            Original sequence.
        position : int
            1-based position in the sequence.
        mutant : str
            Mutant residue.

        Returns
        -------
        str
            Mutated sequence.
        """
        return sequence[: position - 1] + mutant + sequence[position:]

    def apply_mutation(self, mutation_string: str) -> List[str]:
        """Apply multiple mutations to create a mutated sequence.

        Parameters
        ----------
        mutation_strings : str
            List of mutations in format "Chain.Wildtype.Position.Mutant".

        Returns
        -------
        List[str]
            List of mutated sequences (one per chain).

        """
        # Group mutations by chain for efficient processing
        mutations_by_chain: Dict[int, List[Tuple[int, str]]] = {}

        print(f"mutation_string: {mutation_string}")

        if "," in mutation_string:
            individual_mutations = [m.strip() for m in mutation_string.split(",")]
        else:
            individual_mutations = [mutation_string]
        print(f"individual_mutations: {individual_mutations}")
        for single_mut in individual_mutations:
            chain, wildtype, position, mutant = parse_mutation(single_mut)
            chain_idx = get_chain_index(chain)
            if chain_idx not in mutations_by_chain:
                mutations_by_chain[chain_idx] = []
            mutations_by_chain[chain_idx].append((position, mutant))

        # Sort mutations by position (descending) to apply from end to start
        # This avoids position shifting issues
        for chain_idx in mutations_by_chain:
            mutations_by_chain[chain_idx].sort(reverse=True)

        # Create mutated sequences
        mutated_sequence = deepcopy(self.protein_sequence)

        # Apply mutations chain by chain
        for chain_idx, mutations in mutations_by_chain.items():
            seq = mutated_sequence[chain_idx]
            # Apply from end to start to preserve positions
            for position, mutant in mutations:
                seq = self._apply_mutation_to_sequence(seq, position, mutant)
            mutated_sequence[chain_idx] = seq

        return mutated_sequence

    def _expand_single_saturation_mutation(
        self,
        mutation_str: str,
        residue_list: List[str] = STANDARD_AMINO_ACIDS,
    ) -> List[str]:
        """Expand a single saturation mutation.

        Parameters
        ----------
        mutation_str : str
            Saturation mutation in format "Chain.Wildtype.Position".
        residue_list : List[str]
            List of residues to use for saturation.

        Returns
        -------
        List[str]
            List of expanded mutation strings.
        """
        parts = mutation_str.split(".")
        if len(parts) != 3:
            raise ValueError(f"Invalid saturation mutation format: {mutation_str}")

        chain, wildtype, position = parts
        wildtype_upper = wildtype.upper()

        return [
            f"{chain}.{wildtype}.{position}.{res.upper()}"
            for res in residue_list
            if res.upper() != wildtype_upper
        ]

    def expand_saturation_mutations(
        self,
        mutation_strings: List[str],
        residue_list: List[str] = STANDARD_AMINO_ACIDS,
    ) -> List[str]:
        """Expand saturation mutations (mutations without mutant specified).

        A saturation mutation is in format "Chain.Wildtype.Position" (3 parts).
        It will be expanded to all possible mutations using residues in residue_list.
        Handles both single mutations and comma-separated multiple mutations.

        Parameters
        ----------
        mutation_strings : List[str]
            List of mutations, some may be saturation mutations.
        residue_list : List[str]
            List of residues to use for saturation.

        Returns
        -------
        List[str]
            List of expanded mutation strings.
        """
        expanded_mutations = []

        for mut_line in mutation_strings:
            # Handle comma-separated mutations
            if "," in mut_line:
                mutations = [m.strip() for m in mut_line.split(",")]

                # Find saturation mutations
                saturation_indices = [
                    i for i, mut in enumerate(mutations) if len(mut.split(".")) == 3
                ]

                if saturation_indices:
                    # Generate all combinations for saturation mutations
                    expanded_sats = []
                    for idx in saturation_indices:
                        expanded_sats.append(
                            self._expand_single_saturation_mutation(mutations[idx], residue_list)
                        )

                    # Generate all combinations
                    for combo in it.product(*expanded_sats):
                        result_mutations = list(mutations)
                        for combo_idx, sat_idx in enumerate(saturation_indices):
                            result_mutations[sat_idx] = combo[combo_idx]
                        expanded_mutations.append(",".join(result_mutations))
                else:
                    # No saturation mutations, just add as-is
                    expanded_mutations.append(mut_line)
            else:
                # Single mutation
                parts = mut_line.split(".")

                if len(parts) == 3:
                    # Saturation mutation: Chain.Wildtype.Position
                    expanded_mutations.extend(
                        self._expand_single_saturation_mutation(mut_line, residue_list)
                    )
                elif len(parts) == 4:
                    # Regular mutation: Chain.Wildtype.Position.Mutant
                    expanded_mutations.append(mut_line)
                else:
                    raise ValueError(f"Invalid mutation format: {mut_line}")

        return expanded_mutations

    def generate_fasta(
        self,
        mutation_name: str,
        mutated_sequences: List[str],
        output_dir: Union[Path, str],
    ) -> Path:
        """Generate a FASTA file for mutated sequences.

        Parameters
        ----------
        mutation_name : str
            Name/identifier for the mutation.
        mutated_sequences : List[str]
            List of mutated sequences.
        output_dir : Path
            Directory to save FASTA file.

        Returns
        -------
        Path
            Path to the generated FASTA file.
        """
        output_dir = Path(output_dir)
        ensure_directory(output_dir)

        # Format sequences with ':' separator for multi-chain
        sequence_str = ":".join(seq.strip() for seq in mutated_sequences)

        fasta_path = output_dir / f"{mutation_name}.fasta"
        with open(fasta_path, "w") as f:
            f.write(f">{mutation_name}\n{sequence_str}\n")

        return fasta_path

    def generate_all_fastas(
        self,
        output_dir: Union[Path, str],
        mutation_list: List[str],
        residue_list: Optional[List[str]] = None,
    ) -> List[Path]:
        """Generate FASTA files for all mutations.

        Parameters
        ----------
        mutation_list : List[str]
            List of mutation strings (may include saturation mutations).
        output_dir : Path
            Directory to save FASTA files.
        residue_list : Optional[List[str]], optional
            Optional list of residues for saturation expansion, by default None.

        Returns
        -------
        List[Path]
            List of paths to generated FASTA files.
        """
        # Expand saturation mutations if residue list provided
        if residue_list:
            expanded_mutations = self.expand_saturation_mutations(mutation_list, residue_list)
        else:
            expanded_mutations = mutation_list

        fasta_paths = []
        ensure_directory(output_dir)

        for mut_str in expanded_mutations:
            # Apply mutations to generate mutated sequences
            mutated_sequences = self.apply_mutation(mut_str)
            if self.protein_name != "":
                mutation_name = self.protein_name + "-" + mut_str.replace(",", "_")
            else:
                mutation_name = mut_str.replace(",", "_")

            fasta_path = self.generate_fasta(mutation_name, mutated_sequences, output_dir)
            fasta_paths.append(fasta_path)

        return fasta_paths
